


'''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


_lowercase      :      Union[str, Any]            =					{
    """distilbert""": (DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer),
    """roberta""": (RobertaConfig, RobertaForMaskedLM, RobertaTokenizer),
    """bert""": (BertConfig, BertForMaskedLM, BertTokenizer),
    """gpt2""": (GPTaConfig, GPTaLMHeadModel, GPTaTokenizer),
}


def 				lowerCamelCase__				(	A       :      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__				(	A       :      str				,       A       :      Tuple					):




						'''simple docstring'''






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


def 				lowerCamelCase__				(	A       :      Optional[Any]				,       A       :      Dict					):




						'''simple docstring'''






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


def 				lowerCamelCase__				(	):




						'''simple docstring'''






						UpperCAmelCase					    =     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=__lowerCamelCase				,       required=__lowerCamelCase				,       help='''The output directory (log, checkpoints, parameters, etc.)'''					)
						parser.add_argument(
						    '''--data_file'''				,       type=__lowerCamelCase				,       required=__lowerCamelCase				,       help='''The binarized file (tokenized + tokens_to_ids) and grouped by sequence.'''				,       )

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

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

						parser.add_argument('''--temperature'''				,       default=2.0				,       type=__lowerCamelCase				,       help='''Temperature for the softmax temperature.'''					)
						parser.add_argument(
						    '''--alpha_ce'''				,       default=0.5				,       type=__lowerCamelCase				,       help='''Linear weight for the distillation loss. Must be >=0.'''					)
						parser.add_argument(
						    '''--alpha_mlm'''				,       default=0.0				,       type=__lowerCamelCase				,       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=__lowerCamelCase				,       help='''Linear weight for the CLM loss. Must be >=0.'''					)
						parser.add_argument('''--alpha_mse'''				,       default=0.0				,       type=__lowerCamelCase				,       help='''Linear weight of the MSE loss. Must be >=0.'''					)
						parser.add_argument(
						    '''--alpha_cos'''				,       default=0.0				,       type=__lowerCamelCase				,       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=__lowerCamelCase				,       help='''Proportion of tokens for which we need to make a prediction.'''				,       )
						parser.add_argument('''--word_mask'''				,       default=0.8				,       type=__lowerCamelCase				,       help='''Proportion of tokens to mask out.'''					)
						parser.add_argument('''--word_keep'''				,       default=0.1				,       type=__lowerCamelCase				,       help='''Proportion of tokens to keep.'''					)
						parser.add_argument('''--word_rand'''				,       default=0.1				,       type=__lowerCamelCase				,       help='''Proportion of tokens to randomly replace.'''					)
						parser.add_argument(
						    '''--mlm_smoothing'''				,       default=0.7				,       type=__lowerCamelCase				,       help='''Smoothing parameter to emphasize more rare tokens (see XLM, similar to word2vec).'''				,       )
						parser.add_argument('''--token_counts'''				,       type=__lowerCamelCase				,       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=__lowerCamelCase				,       default=3				,       help='''Number of pass on the whole dataset.'''					)
						parser.add_argument('''--batch_size'''				,       type=__lowerCamelCase				,       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=__lowerCamelCase				,       default=50				,       help='''Gradient accumulation for larger training batches.'''				,       )
						parser.add_argument('''--warmup_prop'''				,       default=0.05				,       type=__lowerCamelCase				,       help='''Linear warmup proportion.'''					)
						parser.add_argument('''--weight_decay'''				,       default=0.0				,       type=__lowerCamelCase				,       help='''Weight decay if we apply some.'''					)
						parser.add_argument('''--learning_rate'''				,       default=5E-4				,       type=__lowerCamelCase				,       help='''The initial learning rate for Adam.'''					)
						parser.add_argument('''--adam_epsilon'''				,       default=1E-6				,       type=__lowerCamelCase				,       help='''Epsilon for Adam optimizer.'''					)
						parser.add_argument('''--max_grad_norm'''				,       default=5.0				,       type=__lowerCamelCase				,       help='''Max gradient norm.'''					)
						parser.add_argument('''--initializer_range'''				,       default=0.02				,       type=__lowerCamelCase				,       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=__lowerCamelCase				,       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=__lowerCamelCase				,       default=1				,       help='''Number of GPUs in the node.'''					)
						parser.add_argument('''--local_rank'''				,       type=__lowerCamelCase				,       default=-1				,       help='''Distributed training - Local rank'''					)
						parser.add_argument('''--seed'''				,       type=__lowerCamelCase				,       default=56				,       help='''Random seed'''					)

						parser.add_argument('''--log_interval'''				,       type=__lowerCamelCase				,       default=5_00				,       help='''Tensorboard logging interval.'''					)
						parser.add_argument('''--checkpoint_interval'''				,       type=__lowerCamelCase				,       default=40_00				,       help='''Checkpoint interval.'''					)
						UpperCAmelCase					    =     parser.parse_args()
						sanity_checks(__lowerCamelCase					)

						# ARGS #
						init_gpu_params(__lowerCamelCase					)
						set_seed(__lowerCamelCase					)
						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(__lowerCamelCase					)				,       __lowerCamelCase				,       indent=4					)
												git_log(args.dump_path					)

						UpperCAmelCase					    =     MODEL_CLASSES[args.student_type]
						UpperCAmelCase					    =     MODEL_CLASSES[args.teacher_type]

						# TOKENIZER #
						UpperCAmelCase					    =     teacher_tokenizer_class.from_pretrained(args.teacher_name					)
						UpperCAmelCase					    =     {}
						for tok_name, tok_symbol in tokenizer.special_tokens_map.items():
												UpperCAmelCase					    =     tokenizer.all_special_tokens.index(__lowerCamelCase					)
												UpperCAmelCase					    =     tokenizer.all_special_ids[idx]
						logger.info(f"""Special tokens {special_tok_ids}"""					)
						UpperCAmelCase					    =     special_tok_ids
						UpperCAmelCase					    =     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:
												UpperCAmelCase					    =     pickle.load(__lowerCamelCase					)

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

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

						UpperCAmelCase					    =     LmSeqsDataset(params=__lowerCamelCase				,       data=__lowerCamelCase					)
						logger.info('''Data loader created.'''					)

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

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

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

						# TEACHER #
						UpperCAmelCase					    =     teacher_model_class.from_pretrained(args.teacher_name				,       output_hidden_states=__lowerCamelCase					)
						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(__lowerCamelCase				,       __lowerCamelCase					)
						if args.freeze_token_type_embds:
												freeze_token_type_embeddings(__lowerCamelCase				,       __lowerCamelCase					)

						# 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()
						UpperCAmelCase					    =     Distiller(
						    params=__lowerCamelCase				,       dataset=__lowerCamelCase				,       token_probs=__lowerCamelCase				,       student=__lowerCamelCase				,       teacher=__lowerCamelCase					)
						distiller.train()
						logger.info('''Let\'s go get some drinks.'''					)


if __name__ == "__main__":
							main()









'''simple docstring'''



import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
	import torch

	from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin






class 	UpperCamelCase__:




				def __init__(		self      :			List[str]				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			Optional[int]=None				,   lowerCAmelCase      :			Optional[int]=None				,   lowerCAmelCase      :			List[str]=None				,   lowerCAmelCase      :			int="resnet50"				,   lowerCAmelCase      :			Optional[int]=3				,   lowerCAmelCase      :			Union[str, Any]=32				,   lowerCAmelCase      :			List[Any]=3				,   lowerCAmelCase      :			Tuple=True				,   lowerCAmelCase      :			str=True				,   )->					Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase					    =     parent
										UpperCAmelCase					    =     out_indices if out_indices is not None else [4]
										UpperCAmelCase					    =     stage_names
										UpperCAmelCase					    =     out_features
										UpperCAmelCase					    =     backbone
										UpperCAmelCase					    =     batch_size
										UpperCAmelCase					    =     image_size
										UpperCAmelCase					    =     num_channels
										UpperCAmelCase					    =     use_pretrained_backbone
										UpperCAmelCase					    =     is_training




				def 					a__(		self      :			Optional[Any]				)->					Optional[int]:

										"""simple docstring"""




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

										return config, pixel_values




				def 					a__(		self      :			int				)->					Optional[Any]:

										"""simple docstring"""




										return TimmBackboneConfig(
										    image_size=self.image_size				,   num_channels=self.num_channels				,   out_features=self.out_features				,   out_indices=self.out_indices				,   stage_names=self.stage_names				,   use_pretrained_backbone=self.use_pretrained_backbone				,   backbone=self.backbone				,   )




				def 					a__(		self      :			Dict				,   lowerCAmelCase      :			str				,   lowerCAmelCase      :			Dict				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase					    =     TimmBackbone(config=lowerCAmelCase				)
										model.to(lowerCAmelCase				)
										model.eval()
										with torch.no_grad():
																UpperCAmelCase					    =     model(lowerCAmelCase				)
										self.parent.assertEqual(
										    result.feature_map[-1].shape				,   (self.batch_size, model.channels[-1], 14, 14)				,   )





				def 					a__(		self      :			str				)->					Optional[Any]:

										"""simple docstring"""




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







@require_torch
@require_timm
class 	UpperCamelCase__(       lowerCAmelCase						,    lowerCAmelCase						,    lowerCAmelCase						,    unittest.TestCase	):
				__magic_name__    :				List[str]  		=	(TimmBackbone,) if is_torch_available() else ()
				__magic_name__    :				Any  		=	{"feature-extraction": TimmBackbone} if is_torch_available() else {}
				__magic_name__    :				Union[str, Any]  		=	False
				__magic_name__    :				int  		=	False
				__magic_name__    :				Tuple  		=	False
				__magic_name__    :				List[str]  		=	False




				def 					a__(		self      :			int				)->					str:

										"""simple docstring"""




										UpperCAmelCase					    =     TimmBackboneModelTester(self				)
										UpperCAmelCase					    =     ConfigTester(self				,   config_class=lowerCAmelCase				,   has_text_modality=lowerCAmelCase				)




				def 					a__(		self      :			List[str]				)->					Dict:

										"""simple docstring"""




										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 					a__(		self      :			List[Any]				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase					    =     '''resnet18'''
										UpperCAmelCase					    =     '''microsoft/resnet-18'''

										UpperCAmelCase					    =     AutoBackbone.from_pretrained(lowerCAmelCase				,   use_timm_backbone=lowerCAmelCase				)
										UpperCAmelCase					    =     AutoBackbone.from_pretrained(lowerCAmelCase				)

										self.assertEqual(len(timm_model.out_features				)				,   len(transformers_model.out_features				)				)
										self.assertEqual(len(timm_model.stage_names				)				,   len(transformers_model.stage_names				)				)
										self.assertEqual(timm_model.channels				,   transformers_model.channels				)
										# Out indices are set to the last layer by default. For timm models, we don't know
										# the number of layers in advance, so we set it to (-1,), whereas for transformers
										# models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
										self.assertEqual(timm_model.out_indices				,   (-1,)				)
										self.assertEqual(transformers_model.out_indices				,   [len(timm_model.stage_names				) - 1]				)

										UpperCAmelCase					    =     AutoBackbone.from_pretrained(lowerCAmelCase				,   use_timm_backbone=lowerCAmelCase				,   out_indices=[1, 2, 3]				)
										UpperCAmelCase					    =     AutoBackbone.from_pretrained(lowerCAmelCase				,   out_indices=[1, 2, 3]				)

										self.assertEqual(timm_model.out_indices				,   transformers_model.out_indices				)
										self.assertEqual(len(timm_model.out_features				)				,   len(transformers_model.out_features				)				)
										self.assertEqual(timm_model.channels				,   transformers_model.channels				)




				@unittest.skip('''TimmBackbone doesn\'t support feed forward chunking'''				)
				def 					a__(		self      :			Union[str, Any]				)->					Tuple:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone doesn\'t have num_hidden_layers attribute'''				)
				def 					a__(		self      :			Dict				)->					Dict:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone initialization is managed on the timm side'''				)
				def 					a__(		self      :			Tuple				)->					List[Any]:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds'''				)
				def 					a__(		self      :			List[str]				)->					Optional[Any]:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds'''				)
				def 					a__(		self      :			Optional[Any]				)->					Optional[Any]:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone model cannot be created without specifying a backbone checkpoint'''				)
				def 					a__(		self      :			Any				)->					Optional[int]:

										"""simple docstring"""




										pass




				@unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
				def 					a__(		self      :			Optional[int]				)->					Optional[int]:

										"""simple docstring"""




										pass




				@unittest.skip('''model weights aren\'t tied in TimmBackbone.'''				)
				def 					a__(		self      :			str				)->					str:

										"""simple docstring"""




										pass




				@unittest.skip('''model weights aren\'t tied in TimmBackbone.'''				)
				def 					a__(		self      :			Any				)->					int:

										"""simple docstring"""




										pass




				@unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
				def 					a__(		self      :			List[Any]				)->					int:

										"""simple docstring"""




										pass




				@unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
				def 					a__(		self      :			List[Any]				)->					List[str]:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone doesn\'t have hidden size info in its configuration.'''				)
				def 					a__(		self      :			List[str]				)->					Dict:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone doesn\'t support output_attentions.'''				)
				def 					a__(		self      :			str				)->					Tuple:

										"""simple docstring"""




										pass




				@unittest.skip('''Safetensors is not supported by timm.'''				)
				def 					a__(		self      :			List[Any]				)->					List[Any]:

										"""simple docstring"""




										pass




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

										"""simple docstring"""




										pass




				def 					a__(		self      :			Dict				)->					Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()

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

																UpperCAmelCase					    =     ['''pixel_values''']
																self.assertListEqual(arg_names[:1]				,   lowerCAmelCase				)




				def 					a__(		self      :			Union[str, Any]				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()
										UpperCAmelCase					    =     True
										UpperCAmelCase					    =     self.has_attentions

										# no need to test all models as different heads yield the same functionality
										UpperCAmelCase					    =     self.all_model_classes[0]
										UpperCAmelCase					    =     model_class(lowerCAmelCase				)
										model.to(lowerCAmelCase				)

										UpperCAmelCase					    =     self._prepare_for_class(lowerCAmelCase				,   lowerCAmelCase				)
										UpperCAmelCase					    =     model(**lowerCAmelCase				)
										UpperCAmelCase					    =     outputs[0][-1]

										# Encoder-/Decoder-only models
										UpperCAmelCase					    =     outputs.hidden_states[0]
										hidden_states.retain_grad()

										if self.has_attentions:
																UpperCAmelCase					    =     outputs.attentions[0]
																attentions.retain_grad()

										output.flatten()[0].backward(retain_graph=lowerCAmelCase				)

										self.assertIsNotNone(hidden_states.grad				)

										if self.has_attentions:
																self.assertIsNotNone(attentions.grad				)





				def 					a__(		self      :			Optional[int]				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)
																model.to(lowerCAmelCase				)
																model.eval()
																UpperCAmelCase					    =     model(**lowerCAmelCase				)

																self.assertEqual(len(result.feature_maps				)				,   len(config.out_indices				)				)
																self.assertEqual(len(model.channels				)				,   len(config.out_indices				)				)

																# Check output of last stage is taken if out_features=None, out_indices=None
																UpperCAmelCase					    =     copy.deepcopy(lowerCAmelCase				)
																UpperCAmelCase					    =     None
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)
																model.to(lowerCAmelCase				)
																model.eval()
																UpperCAmelCase					    =     model(**lowerCAmelCase				)

																self.assertEqual(len(result.feature_maps				)				,   1				)
																self.assertEqual(len(model.channels				)				,   1				)

																# Check backbone can be initialized with fresh weights
																UpperCAmelCase					    =     copy.deepcopy(lowerCAmelCase				)
																UpperCAmelCase					    =     False
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)
																model.to(lowerCAmelCase				)
																model.eval()
																UpperCAmelCase					    =     model(**lowerCAmelCase				)





