

import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule







class 					__lowercase  (						_SCREAMING_SNAKE_CASE				):





							"""simple docstring"""

							UpperCAmelCase_						:		BigBirdConfig
							UpperCAmelCase_						:		jnp.dtype								=	jnp.floataa
							UpperCAmelCase_						:		bool								=	True




							def       snake_case			(					self						)		->       Tuple:
														super().setup()
														A      :					List[str]      		=					nn.Dense(5    ,						dtype=self.dtype						)






							def __call__(					self    ,						*__UpperCAmelCase    ,						**__UpperCAmelCase						)		->       int:
														A      :					List[str]      		=					super().__call__(*__UpperCAmelCase    ,						**__UpperCAmelCase						)
														A      :					List[str]      		=					self.cls(outputs[2]						)
														return outputs[:2] + (cls_out,)







class 					__lowercase  (						_SCREAMING_SNAKE_CASE				):





							"""simple docstring"""

							UpperCAmelCase_						:		Optional[int]								=	FlaxBigBirdForNaturalQuestionsModule






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		):
							def cross_entropy(lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_=None		):
														A      :					str      		=					logits.shape[-1]
														A      :					List[str]      		=					(labels[..., None] == jnp.arange(lowerCamelCase_		)[None]).astype('''f4'''		)
														A      :					Optional[int]      		=					jax.nn.log_softmax(lowerCamelCase_     ,				axis=-1		)
														A      :					str      		=					-jnp.sum(labels * logits     ,				axis=-1		)
														if reduction is not None:
																					A      :					int      		=					reduction(lowerCamelCase_		)
														return loss

							A      :					Dict      		=					partial(lowerCamelCase_     ,				reduction=jnp.mean		)
							A      :					List[str]      		=					cross_entropy(lowerCamelCase_     ,				lowerCamelCase_		)
							A      :					List[Any]      		=					cross_entropy(lowerCamelCase_     ,				lowerCamelCase_		)
							A      :					Optional[Any]      		=					cross_entropy(lowerCamelCase_     ,				lowerCamelCase_		)
							return (start_loss + end_loss + pooled_loss) / 3







@dataclass
class 					__lowercase  :





							"""simple docstring"""

							UpperCAmelCase_						:		str								=	"google/bigbird-roberta-base"
							UpperCAmelCase_						:		int								=	30_00
							UpperCAmelCase_						:		int								=	1_05_00

							UpperCAmelCase_						:		int								=	1_28
							UpperCAmelCase_						:		int								=	3

							UpperCAmelCase_						:		int								=	1
							UpperCAmelCase_						:		int								=	5

							# tx_args
							UpperCAmelCase_						:		float								=	3e-5
							UpperCAmelCase_						:		float								=	0.0
							UpperCAmelCase_						:		int								=	2_00_00
							UpperCAmelCase_						:		float								=	0.0095

							UpperCAmelCase_						:		str								=	"bigbird-roberta-natural-questions"
							UpperCAmelCase_						:		str								=	"training-expt"
							UpperCAmelCase_						:		str								=	"data/nq-training.jsonl"
							UpperCAmelCase_						:		str								=	"data/nq-validation.jsonl"




							def       snake_case			(					self						)		->       Tuple:
														os.makedirs(self.base_dir    ,						exist_ok=__UpperCAmelCase						)
														A      :					Dict      		=					os.path.join(self.base_dir    ,						self.save_dir						)
														A      :					Optional[Any]      		=					self.batch_size_per_device * jax.device_count()







@dataclass
class 					__lowercase  :





							"""simple docstring"""

							UpperCAmelCase_						:		int
							UpperCAmelCase_						:		int								=	40_96  # no dynamic padding on TPUs




							def __call__(					self    ,						__UpperCAmelCase						)		->       int:
														A      :					Union[str, Any]      		=					self.collate_fn(__UpperCAmelCase						)
														A      :					Any      		=					jax.tree_util.tree_map(__UpperCAmelCase    ,						__UpperCAmelCase						)
														return batch




							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       int:
														A		,		A      :					Optional[Any]      		=					self.fetch_inputs(features['''input_ids''']						)
														A      :					Optional[int]      		=					{
														    '''input_ids''': jnp.array(__UpperCAmelCase    ,						dtype=jnp.intaa						),
														    '''attention_mask''': jnp.array(__UpperCAmelCase    ,						dtype=jnp.intaa						),
														    '''start_labels''': jnp.array(features['''start_token''']    ,						dtype=jnp.intaa						),
														    '''end_labels''': jnp.array(features['''end_token''']    ,						dtype=jnp.intaa						),
														    '''pooled_labels''': jnp.array(features['''category''']    ,						dtype=jnp.intaa						),
														}
														return batch




							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       Optional[Any]:
														A      :					Tuple      		=					[self._fetch_inputs(__UpperCAmelCase						) for ids in input_ids]
														return zip(*__UpperCAmelCase						)






							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       str:
														A      :					Dict      		=					[1 for _ in range(len(__UpperCAmelCase						)						)]
														while len(__UpperCAmelCase						) < self.max_length:
																					input_ids.append(self.pad_id						)
																					attention_mask.append(0						)
														return input_ids, attention_mask






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_=None		):
							if seed is not None:
														A      :					Optional[Any]      		=					dataset.shuffle(seed=lowerCamelCase_		)
							for i in range(len(lowerCamelCase_		) // batch_size		):
														A      :					List[str]      		=					dataset[i * batch_size : (i + 1) * batch_size]
														yield dict(lowerCamelCase_		)






@partial(jax.pmap     ,				axis_name='''batch'''		)
def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				**lowerCamelCase_		):
							def loss_fn(lowerCamelCase_		):
														A      :					Optional[int]      		=					model_inputs.pop('''start_labels'''		)
														A      :					List[str]      		=					model_inputs.pop('''end_labels'''		)
														A      :					Optional[Any]      		=					model_inputs.pop('''pooled_labels'''		)

														A      :					Optional[Any]      		=					state.apply_fn(**lowerCamelCase_     ,				params=lowerCamelCase_     ,				dropout_rng=lowerCamelCase_     ,				train=lowerCamelCase_		)
														A		,		A		,		A      :					Any      		=					outputs

														return state.loss_fn(
														    lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				)

							A		,		A      :					Dict      		=					jax.random.split(lowerCamelCase_		)
							A      :					str      		=					jax.value_and_grad(lowerCamelCase_		)
							A		,		A      :					int      		=					grad_fn(state.params		)
							A      :					List[Any]      		=					jax.lax.pmean({'''loss''': loss}     ,				axis_name='''batch'''		)
							A      :					Any      		=					jax.lax.pmean(lowerCamelCase_     ,				'''batch'''		)

							A      :					int      		=					state.apply_gradients(grads=lowerCamelCase_		)
							return state, metrics, new_drp_rng






@partial(jax.pmap     ,				axis_name='''batch'''		)
def        snake_case__		(    lowerCamelCase_     ,				**lowerCamelCase_		):
							A      :					Union[str, Any]      		=					model_inputs.pop('''start_labels'''		)
							A      :					Optional[Any]      		=					model_inputs.pop('''end_labels'''		)
							A      :					str      		=					model_inputs.pop('''pooled_labels'''		)

							A      :					Dict      		=					state.apply_fn(**lowerCamelCase_     ,				params=state.params     ,				train=lowerCamelCase_		)
							A		,		A		,		A      :					Optional[int]      		=					outputs

							A      :					Any      		=					state.loss_fn(lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		)
							A      :					Optional[Any]      		=					jax.lax.pmean({'''loss''': loss}     ,				axis_name='''batch'''		)
							return metrics







class 					__lowercase  (						train_state.TrainState				):





							"""simple docstring"""

							UpperCAmelCase_						:		Callable								=	struct.field(pytree_node=_SCREAMING_SNAKE_CASE				)







@dataclass
class 					__lowercase  :





							"""simple docstring"""

							UpperCAmelCase_						:		Args
							UpperCAmelCase_						:		Callable
							UpperCAmelCase_						:		Callable
							UpperCAmelCase_						:		Callable
							UpperCAmelCase_						:		Callable
							UpperCAmelCase_						:		wandb
							UpperCAmelCase_						:		Callable								=	None




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase=None						)		->       List[str]:
														A      :					Optional[Any]      		=					model.params
														A      :					Union[str, Any]      		=					TrainState.create(
														    apply_fn=model.__call__    ,						params=__UpperCAmelCase    ,						tx=__UpperCAmelCase    ,						loss_fn=__UpperCAmelCase    ,						)
														if ckpt_dir is not None:
																					A		,		A		,		A		,		A		,		A      :					List[Any]      		=					restore_checkpoint(__UpperCAmelCase    ,						__UpperCAmelCase						)
																					A      :					List[str]      		=					{
																					    '''lr''': args.lr,
																					    '''init_lr''': args.init_lr,
																					    '''warmup_steps''': args.warmup_steps,
																					    '''num_train_steps''': num_train_steps,
																					    '''weight_decay''': args.weight_decay,
																					}
																					A		,		A      :					str      		=					build_tx(**__UpperCAmelCase						)
																					A      :					List[str]      		=					train_state.TrainState(
																					    step=__UpperCAmelCase    ,						apply_fn=model.__call__    ,						params=__UpperCAmelCase    ,						tx=__UpperCAmelCase    ,						opt_state=__UpperCAmelCase    ,						)
																					A      :					Optional[int]      		=					args
																					A      :					Tuple      		=					data_collator
																					A      :					Any      		=					lr
																					A      :					Dict      		=					params
														A      :					Optional[int]      		=					jax_utils.replicate(__UpperCAmelCase						)
														return state




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)		->       Dict:
														A      :					List[Any]      		=					self.args
														A      :					Tuple      		=					len(__UpperCAmelCase						) // args.batch_size

														A      :					Any      		=					jax.random.PRNGKey(0						)
														A      :					int      		=					jax.random.split(__UpperCAmelCase    ,						jax.device_count()						)
														for epoch in range(args.max_epochs						):
																					A      :					Tuple      		=					jnp.array(0    ,						dtype=jnp.floataa						)
																					A      :					int      		=					get_batched_dataset(__UpperCAmelCase    ,						args.batch_size    ,						seed=__UpperCAmelCase						)
																					A      :					Tuple      		=					0
																					for batch in tqdm(__UpperCAmelCase    ,						total=__UpperCAmelCase    ,						desc=f'Running EPOCH-{epoch}'						):
																												A      :					Any      		=					self.data_collator(__UpperCAmelCase						)
																												A		,		A		,		A      :					Optional[Any]      		=					self.train_step_fn(__UpperCAmelCase    ,						__UpperCAmelCase    ,						**__UpperCAmelCase						)
																												running_loss += jax_utils.unreplicate(metrics['''loss''']						)
																												i += 1
																												if i % args.logging_steps == 0:
																																			A      :					Tuple      		=					jax_utils.unreplicate(state.step						)
																																			A      :					List[Any]      		=					running_loss.item() / i
																																			A      :					Optional[int]      		=					self.scheduler_fn(state_step - 1						)

																																			A      :					str      		=					self.evaluate(__UpperCAmelCase    ,						__UpperCAmelCase						)
																																			A      :					List[str]      		=					{
																																			    '''step''': state_step.item(),
																																			    '''eval_loss''': eval_loss.item(),
																																			    '''tr_loss''': tr_loss,
																																			    '''lr''': lr.item(),
																																			}
																																			tqdm.write(str(__UpperCAmelCase						)						)
																																			self.logger.log(__UpperCAmelCase    ,						commit=__UpperCAmelCase						)

																												if i % args.save_steps == 0:
																																			self.save_checkpoint(args.save_dir + f'-e{epoch}-s{i}'    ,						state=__UpperCAmelCase						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)		->       List[Any]:
														A      :					List[str]      		=					get_batched_dataset(__UpperCAmelCase    ,						self.args.batch_size						)
														A      :					int      		=					len(__UpperCAmelCase						) // self.args.batch_size
														A      :					Optional[int]      		=					jnp.array(0    ,						dtype=jnp.floataa						)
														A      :					Optional[int]      		=					0
														for batch in tqdm(__UpperCAmelCase    ,						total=__UpperCAmelCase    ,						desc='''Evaluating ... '''						):
																					A      :					List[str]      		=					self.data_collator(__UpperCAmelCase						)
																					A      :					List[str]      		=					self.val_step_fn(__UpperCAmelCase    ,						**__UpperCAmelCase						)
																					running_loss += jax_utils.unreplicate(metrics['''loss''']						)
																					i += 1
														return running_loss / i






							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)		->       Any:
														A      :					Optional[int]      		=					jax_utils.unreplicate(__UpperCAmelCase						)
														print(f'SAVING CHECKPOINT IN {save_dir}'    ,						end=''' ... '''						)
														self.model_save_fn(__UpperCAmelCase    ,						params=state.params						)
														with open(os.path.join(__UpperCAmelCase    ,						'''opt_state.msgpack'''						)    ,						'''wb'''						) as f:
																					f.write(to_bytes(state.opt_state						)						)
														joblib.dump(self.args    ,						os.path.join(__UpperCAmelCase    ,						'''args.joblib'''						)						)
														joblib.dump(self.data_collator    ,						os.path.join(__UpperCAmelCase    ,						'''data_collator.joblib'''						)						)
														with open(os.path.join(__UpperCAmelCase    ,						'''training_state.json'''						)    ,						'''w'''						) as f:
																					json.dump({'''step''': state.step.item()}    ,						__UpperCAmelCase						)
														print('''DONE'''						)






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_		):
							print(F'RESTORING CHECKPOINT FROM {save_dir}'     ,				end=''' ... '''		)
							with open(os.path.join(lowerCamelCase_     ,				'''flax_model.msgpack'''		)     ,				'''rb'''		) as f:
														A      :					List[str]      		=					from_bytes(state.params     ,				f.read()		)

							with open(os.path.join(lowerCamelCase_     ,				'''opt_state.msgpack'''		)     ,				'''rb'''		) as f:
														A      :					Optional[int]      		=					from_bytes(state.opt_state     ,				f.read()		)

							A      :					Dict      		=					joblib.load(os.path.join(lowerCamelCase_     ,				'''args.joblib'''		)		)
							A      :					str      		=					joblib.load(os.path.join(lowerCamelCase_     ,				'''data_collator.joblib'''		)		)

							with open(os.path.join(lowerCamelCase_     ,				'''training_state.json'''		)     ,				'''r'''		) as f:
														A      :					Dict      		=					json.load(lowerCamelCase_		)
							A      :					List[Any]      		=					training_state['''step''']

							print('''DONE'''		)
							return params, opt_state, step, args, data_collator






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		):
							A      :					List[str]      		=					num_train_steps - warmup_steps
							A      :					List[Any]      		=					optax.linear_schedule(init_value=lowerCamelCase_     ,				end_value=lowerCamelCase_     ,				transition_steps=lowerCamelCase_		)
							A      :					Optional[Any]      		=					optax.linear_schedule(init_value=lowerCamelCase_     ,				end_value=1E-7     ,				transition_steps=lowerCamelCase_		)
							A      :					Union[str, Any]      		=					optax.join_schedules(schedules=[warmup_fn, decay_fn]     ,				boundaries=[warmup_steps]		)
							return lr






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		):
							def weight_decay_mask(lowerCamelCase_		):
														A      :					int      		=					traverse_util.flatten_dict(lowerCamelCase_		)
														A      :					List[Any]      		=					{k: (v[-1] != '''bias''' and v[-2:] != ('''LayerNorm''', '''scale''')) for k, v in params.items()}
														return traverse_util.unflatten_dict(lowerCamelCase_		)

							A      :					str      		=					scheduler_fn(lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		)

							A      :					str      		=					optax.adamw(learning_rate=lowerCamelCase_     ,				weight_decay=lowerCamelCase_     ,				mask=lowerCamelCase_		)
							return tx, lr









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 					__lowercase  :





							"""simple docstring"""

							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[Any]:
														A      :					Tuple      		=					parent
														A      :					Any      		=					batch_size
														A      :					List[Any]      		=					image_size
														A      :					Tuple      		=					num_channels
														A      :					Any      		=					num_stages
														A      :					Any      		=					hidden_sizes
														A      :					List[str]      		=					depths
														A      :					str      		=					is_training
														A      :					Any      		=					use_labels
														A      :					Any      		=					intermediate_size
														A      :					List[str]      		=					hidden_act
														A      :					List[Any]      		=					type_sequence_label_size
														A      :					Optional[int]      		=					initializer_range
														A      :					Tuple      		=					out_features
														A      :					Tuple      		=					num_labels
														A      :					Tuple      		=					scope
														A      :					int      		=					num_stages




							def       snake_case			(					self						)		->       Optional[Any]:
														A      :					Any      		=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

														A      :					Any      		=					None
														if self.use_labels:
																					A      :					Union[str, Any]      		=					ids_tensor([self.batch_size]    ,						self.type_sequence_label_size						)

														A      :					Any      		=					self.get_config()

														return config, pixel_values, labels




							def       snake_case			(					self						)		->       Optional[int]:
														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						)		->       List[Any]:
														A      :					List[Any]      		=					UperNetForSemanticSegmentation(config=__UpperCAmelCase						)
														model.to(__UpperCAmelCase						)
														model.eval()
														A      :					int      		=					model(__UpperCAmelCase						)
														self.parent.assertEqual(
														    result.logits.shape    ,						(self.batch_size, self.num_labels, self.image_size, self.image_size)						)






							def       snake_case			(					self						)		->       Tuple:
														A      :					List[str]      		=					self.prepare_config_and_inputs()
														(
														    (
														    A
														)		,		(
														    A
														)		,		(
														    A
														)		,		
														)      :					Tuple      		=					config_and_inputs
														A      :					Dict      		=					{'''pixel_values''': pixel_values}
														return config, inputs_dict







@require_torch
class 					__lowercase  (						_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE		,       unittest.TestCase				):





							"""simple docstring"""

							UpperCAmelCase_						:		str								=	(UperNetForSemanticSegmentation,) if is_torch_available() else ()
							UpperCAmelCase_						:		Union[str, Any]								=	{'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
							UpperCAmelCase_						:		int								=	False
							UpperCAmelCase_						:		List[Any]								=	False
							UpperCAmelCase_						:		Tuple								=	False
							UpperCAmelCase_						:		List[Any]								=	False
							UpperCAmelCase_						:		int								=	False
							UpperCAmelCase_						:		Dict								=	False




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




							def       snake_case			(					self						)		->       Optional[int]:
														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						)		->       Optional[int]:
														return




							def       snake_case			(					self						)		->       Any:
														A		,		A      :					Any      		=					self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					A      :					int      		=					model_class(__UpperCAmelCase						)
																					A      :					List[Any]      		=					inspect.signature(model.forward						)
																					# signature.parameters is an OrderedDict => so arg_names order is deterministic
																					A      :					Union[str, Any]      		=					[*signature.parameters.keys()]

																					A      :					Any      		=					['''pixel_values''']
																					self.assertListEqual(arg_names[:1]    ,						__UpperCAmelCase						)




							def       snake_case			(					self						)		->       List[str]:
														A      :					Optional[int]      		=					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[Any]:
														pass




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




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




							@unittest.skip(reason='''UperNet does not have a base model'''						)
							def       snake_case			(					self						)		->       Any:
														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						)		->       List[Any]:
														pass




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




							def       snake_case			(					self						)		->       Union[str, Any]:
														def check_hidden_states_output(__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						):
																					A      :					Any      		=					model_class(__UpperCAmelCase						)
																					model.to(__UpperCAmelCase						)
																					model.eval()

																					with torch.no_grad():
																												A      :					List[str]      		=					model(**self._prepare_for_class(__UpperCAmelCase    ,						__UpperCAmelCase						)						)

																					A      :					List[Any]      		=					outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																					A      :					List[str]      		=					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      :					Optional[Any]      		=					self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					A      :					List[str]      		=					True
																					check_hidden_states_output(__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)

																					# check that output_hidden_states also work using config
																					del inputs_dict["output_hidden_states"]
																					A      :					Any      		=					True

																					check_hidden_states_output(__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)




							def       snake_case			(					self						)		->       Tuple:
														A		,		A      :					Dict      		=					self.model_tester.prepare_config_and_inputs_for_common()

														A      :					Dict      		=					_config_zero_init(__UpperCAmelCase						)
														A      :					List[str]      		=					_config_zero_init(configs_no_init.backbone_config						)
														for model_class in self.all_model_classes:
																					A      :					Any      		=					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						)		->       Tuple:
														pass






							@slow
							def       snake_case			(					self						)		->       str:
														for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					A      :					Tuple      		=					UperNetForSemanticSegmentation.from_pretrained(__UpperCAmelCase						)
																					self.assertIsNotNone(__UpperCAmelCase						)






def        snake_case__		(    ):
							A      :					Tuple      		=					hf_hub_download(
							    repo_id='''hf-internal-testing/fixtures_ade20k'''     ,				repo_type='''dataset'''     ,				filename='''ADE_val_00000001.jpg'''		)
							A      :					Optional[Any]      		=					Image.open(lowerCamelCase_		).convert('''RGB'''		)
							return image







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





							"""simple docstring"""

							def       snake_case			(					self						)		->       Dict:
														A      :					List[Any]      		=					AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny'''						)
														A      :					List[str]      		=					UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny'''						).to(__UpperCAmelCase						)

														A      :					List[str]      		=					prepare_img()
														A      :					str      		=					processor(images=__UpperCAmelCase    ,						return_tensors='''pt'''						).to(__UpperCAmelCase						)

														with torch.no_grad():
																					A      :					str      		=					model(**__UpperCAmelCase						)

														A      :					Optional[int]      		=					torch.Size((1, model.config.num_labels, 5_12, 5_12)						)
														self.assertEqual(outputs.logits.shape    ,						__UpperCAmelCase						)

														A      :					List[Any]      		=					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						)		->       int:
														A      :					List[str]      		=					AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny'''						)
														A      :					Any      		=					UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny'''						).to(__UpperCAmelCase						)

														A      :					str      		=					prepare_img()
														A      :					List[Any]      		=					processor(images=__UpperCAmelCase    ,						return_tensors='''pt'''						).to(__UpperCAmelCase						)

														with torch.no_grad():
																					A      :					List[Any]      		=					model(**__UpperCAmelCase						)

														A      :					List[str]      		=					torch.Size((1, model.config.num_labels, 5_12, 5_12)						)
														self.assertEqual(outputs.logits.shape    ,						__UpperCAmelCase						)

														A      :					Tuple      		=					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						)						)






