



"""simple docstring"""







import argparse
import json
import os
import pickle
import shutil

import numpy as np
import torch
from distiller import Distiller
from lm_seqs_dataset import LmSeqsDataset

from transformers import (
    BertConfig,
    BertForMaskedLM,
    BertTokenizer,
    DistilBertConfig,
    DistilBertForMaskedLM,
    DistilBertTokenizer,
    GPTaConfig,
    GPTaLMHeadModel,
    GPTaTokenizer,
    RobertaConfig,
    RobertaForMaskedLM,
    RobertaTokenizer,
)
from utils import git_log, init_gpu_params, logger, set_seed


UpperCamelCase     =							{
    """distilbert""": (DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer),
    """roberta""": (RobertaConfig, RobertaForMaskedLM, RobertaTokenizer),
    """bert""": (BertConfig, BertForMaskedLM, BertTokenizer),
    """gpt2""": (GPTaConfig, GPTaLMHeadModel, GPTaTokenizer),
}

def       _lowerCamelCase				(							UpperCAmelCase_ :       int     )	-> List[Any]:






	"""simple docstring"""





	assert (args.mlm and args.alpha_mlm > 0.0) or (not args.mlm and args.alpha_mlm == 0.0)
	assert (args.alpha_mlm > 0.0 and args.alpha_clm == 0.0) or (args.alpha_mlm == 0.0 and args.alpha_clm > 0.0)
	if args.mlm:
		assert os.path.isfile(args.token_counts     )
		assert (args.student_type in ["roberta", "distilbert"]) and (args.teacher_type in ["roberta", "bert"])
	else:
		assert (args.student_type in ["gpt2"]) and (args.teacher_type in ["gpt2"])

	assert args.teacher_type == args.student_type or (
	    args.student_type == "distilbert" and args.teacher_type == "bert"
	)
	assert os.path.isfile(args.student_config     )
	if args.student_pretrained_weights is not None:
		assert os.path.isfile(args.student_pretrained_weights     )

	if args.freeze_token_type_embds:
		assert args.student_type in ["roberta"]

	assert args.alpha_ce >= 0.0
	assert args.alpha_mlm >= 0.0
	assert args.alpha_clm >= 0.0
	assert args.alpha_mse >= 0.0
	assert args.alpha_cos >= 0.0
	assert args.alpha_ce + args.alpha_mlm + args.alpha_clm + args.alpha_mse + args.alpha_cos > 0.0

def       _lowerCamelCase				(							UpperCAmelCase_ :       List[str],	UpperCAmelCase_ :       str     )	-> List[Any]:






	"""simple docstring"""





	if args.student_type == "roberta":
		A__											=					False
	elif args.student_type == "gpt2":
		A__											=					False

def       _lowerCamelCase				(							UpperCAmelCase_ :       Any,	UpperCAmelCase_ :       List[Any]     )	-> List[str]:






	"""simple docstring"""





	if args.student_type == "roberta":
		A__											=					False

def       _lowerCamelCase				(							)	-> Any:






	"""simple docstring"""





	A__											=					argparse.ArgumentParser(description="Training"     )
	parser.add_argument("--force",	action="store_true",	help="Overwrite dump_path if it already exists."     )

	parser.add_argument(
	    "--dump_path",	type=UpperCAmelCase_,	required=UpperCAmelCase_,	help="The output directory (log, checkpoints, parameters, etc.)"     )
	parser.add_argument(
	    "--data_file",	type=UpperCAmelCase_,	required=UpperCAmelCase_,	help="The binarized file (tokenized + tokens_to_ids) and grouped by sequence.",	)

	parser.add_argument(
	    "--student_type",	type=UpperCAmelCase_,	choices=["distilbert", "roberta", "gpt2"],	required=UpperCAmelCase_,	help="The student type (DistilBERT, RoBERTa).",	)
	parser.add_argument("--student_config",	type=UpperCAmelCase_,	required=UpperCAmelCase_,	help="Path to the student configuration."     )
	parser.add_argument(
	    "--student_pretrained_weights",	default=UpperCAmelCase_,	type=UpperCAmelCase_,	help="Load student initialization checkpoint."     )

	parser.add_argument(
	    "--teacher_type",	choices=["bert", "roberta", "gpt2"],	required=UpperCAmelCase_,	help="Teacher type (BERT, RoBERTa)."     )
	parser.add_argument("--teacher_name",	type=UpperCAmelCase_,	required=UpperCAmelCase_,	help="The teacher model."     )

	parser.add_argument("--temperature",	default=2.0,	type=UpperCAmelCase_,	help="Temperature for the softmax temperature."     )
	parser.add_argument(
	    "--alpha_ce",	default=0.5,	type=UpperCAmelCase_,	help="Linear weight for the distillation loss. Must be >=0."     )
	parser.add_argument(
	    "--alpha_mlm",	default=0.0,	type=UpperCAmelCase_,	help="Linear weight for the MLM loss. Must be >=0. Should be used in conjunction with `mlm` flag.",	)
	parser.add_argument("--alpha_clm",	default=0.5,	type=UpperCAmelCase_,	help="Linear weight for the CLM loss. Must be >=0."     )
	parser.add_argument("--alpha_mse",	default=0.0,	type=UpperCAmelCase_,	help="Linear weight of the MSE loss. Must be >=0."     )
	parser.add_argument(
	    "--alpha_cos",	default=0.0,	type=UpperCAmelCase_,	help="Linear weight of the cosine embedding loss. Must be >=0."     )

	parser.add_argument(
	    "--mlm",	action="store_true",	help="The LM step: MLM or CLM. If `mlm` is True, the MLM is used over CLM."     )
	parser.add_argument(
	    "--mlm_mask_prop",	default=0.15,	type=UpperCAmelCase_,	help="Proportion of tokens for which we need to make a prediction.",	)
	parser.add_argument("--word_mask",	default=0.8,	type=UpperCAmelCase_,	help="Proportion of tokens to mask out."     )
	parser.add_argument("--word_keep",	default=0.1,	type=UpperCAmelCase_,	help="Proportion of tokens to keep."     )
	parser.add_argument("--word_rand",	default=0.1,	type=UpperCAmelCase_,	help="Proportion of tokens to randomly replace."     )
	parser.add_argument(
	    "--mlm_smoothing",	default=0.7,	type=UpperCAmelCase_,	help="Smoothing parameter to emphasize more rare tokens (see XLM, similar to word2vec).",	)
	parser.add_argument("--token_counts",	type=UpperCAmelCase_,	help="The token counts in the data_file for MLM."     )

	parser.add_argument(
	    "--restrict_ce_to_mask",	action="store_true",	help="If true, compute the distillation loss only the [MLM] prediction distribution.",	)
	parser.add_argument(
	    "--freeze_pos_embs",	action="store_true",	help="Freeze positional embeddings during distillation. For student_type in ['roberta', 'gpt2'] only.",	)
	parser.add_argument(
	    "--freeze_token_type_embds",	action="store_true",	help="Freeze token type embeddings during distillation if existent. For student_type in ['roberta'] only.",	)

	parser.add_argument("--n_epoch",	type=UpperCAmelCase_,	default=3,	help="Number of pass on the whole dataset."     )
	parser.add_argument("--batch_size",	type=UpperCAmelCase_,	default=5,	help="Batch size (for each process)."     )
	parser.add_argument(
	    "--group_by_size",	action="store_false",	help="If true, group sequences that have similar length into the same batch. Default is true.",	)

	parser.add_argument(
	    "--gradient_accumulation_steps",	type=UpperCAmelCase_,	default=50,	help="Gradient accumulation for larger training batches.",	)
	parser.add_argument("--warmup_prop",	default=0.05,	type=UpperCAmelCase_,	help="Linear warmup proportion."     )
	parser.add_argument("--weight_decay",	default=0.0,	type=UpperCAmelCase_,	help="Weight decay if we apply some."     )
	parser.add_argument("--learning_rate",	default=5e-4,	type=UpperCAmelCase_,	help="The initial learning rate for Adam."     )
	parser.add_argument("--adam_epsilon",	default=1e-6,	type=UpperCAmelCase_,	help="Epsilon for Adam optimizer."     )
	parser.add_argument("--max_grad_norm",	default=5.0,	type=UpperCAmelCase_,	help="Max gradient norm."     )
	parser.add_argument("--initializer_range",	default=0.02,	type=UpperCAmelCase_,	help="Random initialization range."     )

	parser.add_argument(
	    "--fp16",	action="store_true",	help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",	)
	parser.add_argument(
	    "--fp16_opt_level",	type=UpperCAmelCase_,	default="O1",	help=(
	        "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
	        "See details at https://nvidia.github.io/apex/amp.html"
	    ),	)
	parser.add_argument("--n_gpu",	type=UpperCAmelCase_,	default=1,	help="Number of GPUs in the node."     )
	parser.add_argument("--local_rank",	type=UpperCAmelCase_,	default=-1,	help="Distributed training - Local rank"     )
	parser.add_argument("--seed",	type=UpperCAmelCase_,	default=56,	help="Random seed"     )

	parser.add_argument("--log_interval",	type=UpperCAmelCase_,	default=500,	help="Tensorboard logging interval."     )
	parser.add_argument("--checkpoint_interval",	type=UpperCAmelCase_,	default=4000,	help="Checkpoint interval."     )
	A__											=					parser.parse_args()
	sanity_checks(UpperCAmelCase_     )

	# ARGS #
	init_gpu_params(UpperCAmelCase_     )
	set_seed(UpperCAmelCase_     )
	if args.is_master:
		if os.path.exists(args.dump_path     ):
			if not args.force:
				raise ValueError(
				    F"""Serialization dir {args.dump_path} already exists, but you have not precised wheter to overwrite"""
				    " itUse `--force` if you want to overwrite it"     )
			else:
				shutil.rmtree(args.dump_path     )

		if not os.path.exists(args.dump_path     ):
			os.makedirs(args.dump_path     )
		logger.info(F"""Experiment will be dumped and logged in {args.dump_path}"""     )

		# SAVE PARAMS #
		logger.info(F"""Param: {args}"""     )
		with open(os.path.join(args.dump_path,	"parameters.json"     ),	"w"     ) as f:
			json.dump(vars(UpperCAmelCase_     ),	UpperCAmelCase_,	indent=4     )
		git_log(args.dump_path     )

	A__							,							A__							,							A__											=					MODEL_CLASSES[args.student_type]
	A__							,							A__							,							A__											=					MODEL_CLASSES[args.teacher_type]

	# TOKENIZER #
	A__											=					teacher_tokenizer_class.from_pretrained(args.teacher_name     )
	A__											=					{}
	for tok_name, tok_symbol in tokenizer.special_tokens_map.items():
		A__											=					tokenizer.all_special_tokens.index(UpperCAmelCase_     )
		A__											=					tokenizer.all_special_ids[idx]
	logger.info(F"""Special tokens {special_tok_ids}"""     )
	A__											=					special_tok_ids
	A__											=					tokenizer.max_model_input_sizes[args.teacher_name]

	# DATA LOADER #
	logger.info(F"""Loading data from {args.data_file}"""     )
	with open(args.data_file,	"rb"     ) as fp:
		A__											=					pickle.load(UpperCAmelCase_     )

	if args.mlm:
		logger.info(F"""Loading token counts from {args.token_counts} (already pre-computed)"""     )
		with open(args.token_counts,	"rb"     ) as fp:
			A__											=					pickle.load(UpperCAmelCase_     )

		A__											=					np.maximum(UpperCAmelCase_,	1     ) ** -args.mlm_smoothing
		for idx in special_tok_ids.values():
			A__											=					0.0  # do not predict special tokens
		A__											=					torch.from_numpy(UpperCAmelCase_     )
	else:
		A__											=					None

	A__											=					LmSeqsDataset(params=UpperCAmelCase_,	data=UpperCAmelCase_     )
	logger.info("Data loader created."     )

	# STUDENT #
	logger.info(F"""Loading student config from {args.student_config}"""     )
	A__											=					student_config_class.from_pretrained(args.student_config     )
	A__											=					True

	if args.student_pretrained_weights is not None:
		logger.info(F"""Loading pretrained weights from {args.student_pretrained_weights}"""     )
		A__											=					student_model_class.from_pretrained(args.student_pretrained_weights,	config=UpperCAmelCase_     )
	else:
		A__											=					student_model_class(UpperCAmelCase_     )

	if args.n_gpu > 0:
		student.to(F"""cuda:{args.local_rank}"""     )
	logger.info("Student loaded."     )

	# TEACHER #
	A__											=					teacher_model_class.from_pretrained(args.teacher_name,	output_hidden_states=UpperCAmelCase_     )
	if args.n_gpu > 0:
		teacher.to(F"""cuda:{args.local_rank}"""     )
	logger.info(F"""Teacher loaded from {args.teacher_name}."""     )

	# FREEZING #
	if args.freeze_pos_embs:
		freeze_pos_embeddings(UpperCAmelCase_,	UpperCAmelCase_     )
	if args.freeze_token_type_embds:
		freeze_token_type_embeddings(UpperCAmelCase_,	UpperCAmelCase_     )

	# SANITY CHECKS #
	assert student.config.vocab_size == teacher.config.vocab_size
	assert student.config.hidden_size == teacher.config.hidden_size
	assert student.config.max_position_embeddings == teacher.config.max_position_embeddings
	if args.mlm:
		assert token_probs.size(0     ) == stu_architecture_config.vocab_size

	# DISTILLER #
	torch.cuda.empty_cache()
	A__											=					Distiller(
	    params=UpperCAmelCase_,	dataset=UpperCAmelCase_,	token_probs=UpperCAmelCase_,	student=UpperCAmelCase_,	teacher=UpperCAmelCase_     )
	distiller.train()
	logger.info("Let's go get some drinks."     )


if __name__ == "__main__":
			main()









import inspect
import unittest

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

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


if is_torch_available():
							import torch

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


if is_vision_available():
							from PIL import Image

							from transformers import SegformerImageProcessor


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






					def      _a				(  self       :							List[str]       ):






								"""simple docstring"""


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







class   __lowerCAmelCase  :
					"""simple docstring"""






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






								"""simple docstring"""


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



					def      _a				(  self       :							int       ):






								"""simple docstring"""


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

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

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



					def      _a				(  self       :							int       ):






								"""simple docstring"""


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



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






								"""simple docstring"""


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



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






								"""simple docstring"""


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



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






								"""simple docstring"""


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



					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


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







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






					A__					:		Optional[int]      							=     (
					    (
					        SegformerModel,
					        SegformerForSemanticSegmentation,
					        SegformerForImageClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					A__					:		Union[str, Any]      							=     (
					    {
					        "feature-extraction": SegformerModel,
					        "image-classification": SegformerForImageClassification,
					        "image-segmentation": SegformerForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					A__					:		Optional[Any]      							=     True
					A__					:		str      							=     False
					A__					:		Tuple      							=     False
					A__					:		Dict      							=     False



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






								"""simple docstring"""


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



					def      _a				(  self       :							Optional[int]       ):






								"""simple docstring"""


								self.config_tester.run_common_tests()



					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


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



					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


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



					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


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



					@unittest.skip('SegFormer does not use inputs_embeds'       )
					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


								pass



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






								"""simple docstring"""


								pass



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()

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

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



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


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

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

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

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

											self.assertEqual(len(_snake_case       )					,       _snake_case       )

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

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

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

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

											A__	  =     outputs.attentions

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



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






								"""simple docstring"""


								def check_hidden_states_output(_snake_case       :							Dict					,       _snake_case       :							int					,       _snake_case       :							List[Any]       ):
											A__	  =     model_class(_snake_case       )
											model.to(_snake_case       )
											model.eval()

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

											A__	  =     outputs.hidden_states

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

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

								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()

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

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

											check_hidden_states_output(_snake_case					,       _snake_case					,       _snake_case       )



					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


								if not self.model_tester.is_training:
											return

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

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

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



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






								"""simple docstring"""


								pass



					@slow
					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


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







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






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






					@slow
					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


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

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

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

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

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



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






								"""simple docstring"""


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

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

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

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

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



					@slow
					def      _a				(  self       :							Any       ):






								"""simple docstring"""


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

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

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

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

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

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




