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


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





def 					UpperCAmelCase_		(__a      :    Optional[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 					UpperCAmelCase_		(__a      :    Optional[int] ,	__a      :    Union[str, Any] ):






   """simple docstring"""

   if args.student_type == "roberta":
      _a				:	Optional[Any]			=   False
   elif args.student_type == "gpt2":
      _a				:	int			=   False





def 					UpperCAmelCase_		(__a      :    Any ,	__a      :    Dict ):






   """simple docstring"""

   if args.student_type == "roberta":
      _a				:	Dict			=   False





def 					UpperCAmelCase_		():






   """simple docstring"""

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

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

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

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

   parser.add_argument('--log_interval' ,	type=__a ,	default=5_0_0 ,	help='Tensorboard logging interval.' )
   parser.add_argument('--checkpoint_interval' ,	type=__a ,	default=4_0_0_0 ,	help='Checkpoint interval.' )
   _a				:	Any			=   parser.parse_args()
   sanity_checks(__a )

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

   _a,  _a,  _a				:	Union[str, Any]			=   MODEL_CLASSES[args.student_type]
   _a,  _a,  _a				:	Tuple			=   MODEL_CLASSES[args.teacher_type]

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

   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				:	Tuple			=   pickle.load(__a )

      _a				:	Union[str, Any]			=   np.maximum(__a ,	1 ) ** -args.mlm_smoothing
      for idx in special_tok_ids.values():
         _a				:	str			=   0.0  # do not predict special tokens
      _a				:	Optional[int]			=   torch.from_numpy(__a )
   else:
      _a				:	Optional[Any]			=   None

   _a				:	Union[str, Any]			=   LmSeqsDataset(params=__a ,	data=__a )
   logger.info('Data loader created.' )

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

   if args.student_pretrained_weights is not None:
      logger.info(f"""Loading pretrained weights from {args.student_pretrained_weights}""" )
      _a				:	List[str]			=   student_model_class.from_pretrained(args.student_pretrained_weights ,	config=__a )
   else:
      _a				:	List[str]			=   student_model_class(__a )

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

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

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


if __name__ == "__main__":
  main()





'''simple docstring'''







import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

  if is_vision_available():
    from transformers import MaskaFormerImageProcessor

if is_vision_available():
  from PIL import Image







class UpperCAmelCase__	:





  """simple docstring"""

  def __init__(		self      :     int							,_a      :     Any							,_a      :     Optional[int]=2							,_a      :     Optional[Any]=True							,_a      :     Dict=False							,_a      :     Dict=10							,_a      :     Any=3							,_a      :     str=32 * 8							,_a      :     Optional[int]=32 * 8							,_a      :     int=4							,_a      :     str=64							,):


     '''simple docstring'''






     _a				:	Dict			=   parent
     _a				:	Union[str, Any]			=   batch_size
     _a				:	Tuple			=   is_training
     _a				:	List[str]			=   use_auxiliary_loss
     _a				:	Optional[Any]			=   num_queries
     _a				:	str			=   num_channels
     _a				:	List[str]			=   min_size
     _a				:	int			=   max_size
     _a				:	Optional[int]			=   num_labels
     _a				:	List[str]			=   hidden_dim
     _a				:	int			=   hidden_dim


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


     '''simple docstring'''






     _a				:	Tuple			=   floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]       ).to(
         _a       )

     _a				:	Optional[Any]			=   torch.ones([self.batch_size, self.min_size, self.max_size]							,device=_a       )

     _a				:	Union[str, Any]			=   (
         torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]							,device=_a       ) > 0.5
     ).float()
     _a				:	Tuple			=   (torch.rand((self.batch_size, self.num_labels)							,device=_a       ) > 0.5).long()

     _a				:	Dict			=   self.get_config()
     return config, pixel_values, pixel_mask, mask_labels, class_labels


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


     '''simple docstring'''






     _a				:	int			=   MaskaFormerConfig(
         hidden_size=self.hidden_dim							,)
     _a				:	str			=   self.num_queries
     _a				:	Union[str, Any]			=   self.num_labels

     _a				:	Tuple			=   [1, 1, 1, 1]
     _a				:	Dict			=   self.num_channels

     _a				:	str			=   64
     _a				:	Tuple			=   128
     _a				:	Optional[Any]			=   self.hidden_dim
     _a				:	Union[str, Any]			=   self.hidden_dim
     _a				:	List[Any]			=   self.hidden_dim
     return config


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a,  _a,  _a,  _a,  _a				:	Optional[Any]			=   self.prepare_config_and_inputs()
     _a				:	str			=   {'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
     return config, inputs_dict


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


     '''simple docstring'''






     _a				:	str			=   output.encoder_hidden_states
     _a				:	Any			=   output.pixel_decoder_hidden_states
     _a				:	Optional[Any]			=   output.transformer_decoder_hidden_states

     self.parent.assertTrue(len(_a       )							,len(config.backbone_config.depths       )       )
     self.parent.assertTrue(len(_a       )							,len(config.backbone_config.depths       )       )
     self.parent.assertTrue(len(_a       )							,config.decoder_layers       )


  def   __lowercase		(		self      :     List[str]							,_a      :     str							,_a      :     List[Any]							,_a      :     Any							,_a      :     Union[str, Any]=False       ):


     '''simple docstring'''






     with torch.no_grad():
        _a				:	str			=   MaskaFormerModel(config=_a       )
        model.to(_a       )
        model.eval()

        _a				:	Any			=   model(pixel_values=_a							,pixel_mask=_a       )
        _a				:	Optional[Any]			=   model(_a							,output_hidden_states=_a       )

     self.parent.assertEqual(
         output.transformer_decoder_last_hidden_state.shape							,(self.batch_size, self.num_queries, self.hidden_dim)							,)
     # let's ensure the other two hidden state exists
     self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None       )
     self.parent.assertTrue(output.encoder_last_hidden_state is not None       )

     if output_hidden_states:
        self.check_output_hidden_state(_a							,_a       )


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


     '''simple docstring'''






     _a				:	int			=   MaskaFormerForUniversalSegmentation(config=_a       )
     model.to(_a       )
     model.eval()

     def comm_check_on_output(_a      :     Any       ):
        # let's still check that all the required stuff is there
        self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None       )
        self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None       )
        self.parent.assertTrue(result.encoder_last_hidden_state is not None       )
        # okay, now we need to check the logits shape
        # due to the encoder compression, masks have a //4 spatial size
        self.parent.assertEqual(
            result.masks_queries_logits.shape							,(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)							,)
        # + 1 for null class
        self.parent.assertEqual(
            result.class_queries_logits.shape							,(self.batch_size, self.num_queries, self.num_labels + 1)       )

     with torch.no_grad():
        _a				:	Any			=   model(pixel_values=_a							,pixel_mask=_a       )
        _a				:	Optional[int]			=   model(_a       )

        comm_check_on_output(_a       )

        _a				:	List[str]			=   model(
            pixel_values=_a							,pixel_mask=_a							,mask_labels=_a							,class_labels=_a       )

     comm_check_on_output(_a       )

     self.parent.assertTrue(result.loss is not None       )
     self.parent.assertEqual(result.loss.shape							,torch.Size([1]       )       )







@require_torch
class UpperCAmelCase__	( lowercase__    ,				lowercase__    ,				unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				Optional[int]   =					(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
  __UpperCAmelCase				:				Dict   =					{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

  __UpperCAmelCase				:				Dict   =					False
  __UpperCAmelCase				:				Tuple   =					False
  __UpperCAmelCase				:				Dict   =					False
  __UpperCAmelCase				:				List[Any]   =					False


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   MaskaFormerModelTester(self       )
     _a				:	Dict			=   ConfigTester(self							,config_class=_a							,has_text_modality=_a       )


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     self.config_tester.run_common_tests()


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a,  _a				:	List[str]			=   self.model_tester.prepare_config_and_inputs_for_common()
     self.model_tester.create_and_check_maskaformer_model(_a							,**_a							,output_hidden_states=_a       )


  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     _a				:	Dict			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*_a       )


  @unittest.skip(reason='Mask2Former does not use inputs_embeds'       )
  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     pass


  @unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'       )
  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     pass


  @unittest.skip(reason='Mask2Former is not a generative model'       )
  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     pass


  @unittest.skip(reason='Mask2Former does not use token embeddings'       )
  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     pass


  @require_torch_multi_gpu
  @unittest.skip(
      reason='Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'       )
  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     pass


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


     '''simple docstring'''






     pass


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a,  _a				:	Any			=   self.model_tester.prepare_config_and_inputs_for_common()

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

        _a				:	List[Any]			=   ['pixel_values']
        self.assertListEqual(arg_names[:1]							,_a       )


  @slow
  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
        _a				:	Dict			=   MaskaFormerModel.from_pretrained(_a       )
        self.assertIsNotNone(_a       )


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     _a				:	int			=   (self.model_tester.min_size,) * 2
     _a				:	Any			=   {
         'pixel_values': torch.randn((2, 3, *size)							,device=_a       ),
         'mask_labels': torch.randn((2, 10, *size)							,device=_a       ),
         'class_labels': torch.zeros(2							,10							,device=_a       ).long(),
     }
     _a				:	List[Any]			=   self.model_tester.get_config()

     _a				:	int			=   MaskaFormerForUniversalSegmentation(_a       ).to(_a       )
     _a				:	str			=   model(**_a       )
     self.assertTrue(outputs.loss is not None       )


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a,  _a				:	List[Any]			=   self.model_tester.prepare_config_and_inputs_for_common()
     self.model_tester.create_and_check_maskaformer_model(_a							,**_a							,output_hidden_states=_a       )


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a,  _a				:	Dict			=   self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        _a				:	Any			=   model_class(_a       ).to(_a       )
        _a				:	Optional[int]			=   model(**_a							,output_attentions=_a       )
        self.assertTrue(outputs.attentions is not None       )


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     if not self.model_tester.is_training:
        return

     _a				:	List[str]			=   self.all_model_classes[1]
     _a,  _a,  _a,  _a,  _a				:	List[str]			=   self.model_tester.prepare_config_and_inputs()

     _a				:	Any			=   model_class(_a       )
     model.to(_a       )
     model.train()

     _a				:	Union[str, Any]			=   model(_a							,mask_labels=_a							,class_labels=_a       ).loss
     loss.backward()


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a				:	int			=   self.all_model_classes[1]
     _a,  _a,  _a,  _a,  _a				:	List[Any]			=   self.model_tester.prepare_config_and_inputs()
     _a				:	str			=   True
     _a				:	str			=   True

     _a				:	List[str]			=   model_class(_a       ).to(_a       )
     model.train()

     _a				:	Optional[int]			=   model(_a							,mask_labels=_a							,class_labels=_a       )

     _a				:	Tuple			=   outputs.encoder_hidden_states[0]
     encoder_hidden_states.retain_grad()

     _a				:	str			=   outputs.pixel_decoder_hidden_states[0]
     pixel_decoder_hidden_states.retain_grad()

     _a				:	Dict			=   outputs.transformer_decoder_hidden_states[0]
     transformer_decoder_hidden_states.retain_grad()

     _a				:	List[str]			=   outputs.attentions[0]
     attentions.retain_grad()

     outputs.loss.backward(retain_graph=_a       )

     self.assertIsNotNone(encoder_hidden_states.grad       )
     self.assertIsNotNone(pixel_decoder_hidden_states.grad       )
     self.assertIsNotNone(transformer_decoder_hidden_states.grad       )
     self.assertIsNotNone(attentions.grad       )


__lowerCAmelCase         =    1e-4





def 					UpperCAmelCase_		():






   """simple docstring"""

   _a				:	int			=   Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png' )
   return image







@require_vision
@slow
class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  @cached_property
  def   __lowercase		(		self      :     Union[str, Any]       ):


     '''simple docstring'''






     return "facebook/mask2former-swin-small-coco-instance"


  @cached_property
  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints       ) if is_vision_available() else None


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	List[str]			=   MaskaFormerModel.from_pretrained(self.model_checkpoints       ).to(_a       )
     _a				:	int			=   self.default_image_processor
     _a				:	Tuple			=   prepare_img()
     _a				:	Any			=   image_processor(_a							,return_tensors='pt'       ).to(_a       )
     _a				:	Union[str, Any]			=   inputs['pixel_values'].shape
     # check size is divisible by 32
     self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0       )
     # check size
     self.assertEqual(_a							,(1, 3, 384, 384)       )

     with torch.no_grad():
        _a				:	Optional[Any]			=   model(**_a       )

     _a				:	List[Any]			=   torch.tensor(
         [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]       ).to(_a       )
     self.assertTrue(
         torch.allclose(
             outputs.encoder_last_hidden_state[0, 0, :3, :3]							,_a							,atol=_a       )       )

     _a				:	str			=   torch.tensor(
         [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]       ).to(_a       )
     self.assertTrue(
         torch.allclose(
             outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]							,_a							,atol=_a       )       )

     _a				:	Any			=   torch.tensor(
         [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]       ).to(_a       )
     self.assertTrue(
         torch.allclose(
             outputs.transformer_decoder_last_hidden_state[0, :3, :3]							,_a							,atol=_a       )       )


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     _a				:	List[Any]			=   MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(_a       ).eval()
     _a				:	Optional[Any]			=   self.default_image_processor
     _a				:	List[Any]			=   prepare_img()
     _a				:	str			=   image_processor(_a							,return_tensors='pt'       ).to(_a       )
     _a				:	Any			=   inputs['pixel_values'].shape
     # check size is divisible by 32
     self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0       )
     # check size
     self.assertEqual(_a							,(1, 3, 384, 384)       )

     with torch.no_grad():
        _a				:	Optional[int]			=   model(**_a       )
     # masks_queries_logits
     _a				:	Dict			=   outputs.masks_queries_logits
     self.assertEqual(
         masks_queries_logits.shape							,(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)       )
     _a				:	Dict			=   [
         [-8.7839, -9.0056, -8.8121],
         [-7.4104, -7.0313, -6.5401],
         [-6.6105, -6.3427, -6.4675],
     ]
     _a				:	Optional[Any]			=   torch.tensor(_a       ).to(_a       )
     self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]							,_a							,atol=_a       )       )
     # class_queries_logits
     _a				:	str			=   outputs.class_queries_logits
     self.assertEqual(class_queries_logits.shape							,(1, model.config.num_queries, model.config.num_labels + 1)       )
     _a				:	str			=   torch.tensor(
         [
             [1.8324, -8.0835, -4.1922],
             [0.8450, -9.0050, -3.6053],
             [0.3045, -7.7293, -3.0275],
         ]       ).to(_a       )
     self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]							,_a							,atol=_a       )       )


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a				:	Any			=   MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(_a       ).eval()
     _a				:	Tuple			=   self.default_image_processor

     _a				:	Tuple			=   image_processor(
         [np.zeros((3, 800, 1333)       ), np.zeros((3, 800, 1333)       )]							,segmentation_maps=[np.zeros((384, 384)       ).astype(np.floataa       ), np.zeros((384, 384)       ).astype(np.floataa       )]							,return_tensors='pt'							,)

     _a				:	str			=   inputs['pixel_values'].to(_a       )
     _a				:	str			=   [el.to(_a       ) for el in inputs['mask_labels']]
     _a				:	Dict			=   [el.to(_a       ) for el in inputs['class_labels']]

     with torch.no_grad():
        _a				:	List[str]			=   model(**_a       )

     self.assertTrue(outputs.loss is not None       )



