


"""simple docstring"""







from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
       from fairscale.optim import OSS


lowerCamelCase_       :       Optional[Any]								=  logging.get_logger(__name__)

lowerCamelCase_       :       List[Any]								=  {
    """linear""": get_linear_schedule_with_warmup,
    """cosine""": get_cosine_schedule_with_warmup,
    """cosine_w_restarts""": get_cosine_with_hard_restarts_schedule_with_warmup,
    """polynomial""": get_polynomial_decay_schedule_with_warmup,
    """constant""": get_constant_schedule,
    """constant_w_warmup""": get_constant_schedule_with_warmup,
}






class 				__A    (	_SCREAMING_SNAKE_CASE      ):






 """simple docstring"""




 def __init__(		self			,   __A=None			,   __A=None			,   *__A			,   **__A						) ->					Optional[Any]:
     super().__init__(*__A			,   **__A						)

     if config is None:
         assert isinstance(self.model			,   __A						), (
             "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
             f''' {self.model.__class__}'''
         )
         a					=self.model.config
     else:
         a					=config

     a					=data_args
     a					=self.config.tgt_vocab_size if isinstance(self.config			,   __A						) else self.config.vocab_size

     if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
         assert self.config.pad_token_id is not None, (
             "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
             " calculation or doing label smoothing."
         )

     if self.config.pad_token_id is None and self.config.eos_token_id is not None:
         logger.warning(
             f'''The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'''
             ''' padding..'''						)

     if self.args.label_smoothing == 0:
         a					=torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id						)
     else:
         # dynamically import label_smoothed_nll_loss
         from utils import label_smoothed_nll_loss

         a					=label_smoothed_nll_loss


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A						) ->					List[str]:
     if self.optimizer is None:
         a					=['''bias''', '''LayerNorm.weight''']
         a					=[
             {
                 '''params''': [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay						)],
                 '''weight_decay''': self.args.weight_decay,
             },
             {
                 '''params''': [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay						)],
                 '''weight_decay''': 0.0,
             },
         ]
         a					=Adafactor if self.args.adafactor else AdamW
         if self.args.adafactor:
             a					=Adafactor
             a					={'''scale_parameter''': False, '''relative_step''': False}
         else:
             a					=AdamW
             a					={
                 '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
                 '''eps''': self.args.adam_epsilon,
             }
         a					=self.args.learning_rate
         if self.sharded_ddp:
             a					=OSS(
                 params=__A			,   optim=__A			,   **__A			,   )
         else:
             a					=optimizer_cls(__A			,   **__A						)

     if self.lr_scheduler is None:
         a					=self._get_lr_scheduler(__A						)
     else:  # ignoring --lr_scheduler
         logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''						)


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A						) ->					Tuple:
     a					=arg_to_scheduler[self.args.lr_scheduler]
     if self.args.lr_scheduler == "constant":
         a					=schedule_func(self.optimizer						)
     elif self.args.lr_scheduler == "constant_w_warmup":
         a					=schedule_func(self.optimizer			,   num_warmup_steps=self.args.warmup_steps						)
     else:
         a					=schedule_func(
             self.optimizer			,   num_warmup_steps=self.args.warmup_steps			,   num_training_steps=__A						)
     return scheduler


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[torch.utils.data.Sampler]:
     if isinstance(self.train_dataset			,   torch.utils.data.IterableDataset						):
         return None
     elif is_torch_tpu_available():
         return get_tpu_sampler(self.train_dataset						)
     else:
         if self.args.sortish_sampler:
             self.train_dataset.make_sortish_sampler(
                 self.args.per_device_train_batch_size			,   distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)			,   )

         return (
             RandomSampler(self.train_dataset						)
             if self.args.local_rank == -1
             else DistributedSampler(self.train_dataset						)
         )


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A						) ->					Optional[int]:
     if self.args.label_smoothing == 0:
         if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
             # force training to ignore pad token
             a					=model(**__A			,   use_cache=__A						)[0]
             a					=self.loss_fn(logits.view(-1			,   logits.shape[-1]						)			,   labels.view(-1						)						)
         else:
             # compute usual loss via models
             a		,				a					=model(**__A			,   labels=__A			,   use_cache=__A						)[:2]
     else:
         # compute label smoothed loss
         a					=model(**__A			,   use_cache=__A						)[0]
         a					=torch.nn.functional.log_softmax(__A			,   dim=-1						)
         a		,				a					=self.loss_fn(__A			,   __A			,   self.args.label_smoothing			,   ignore_index=self.config.pad_token_id						)
     return loss, logits


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A						) ->					Optional[Any]:
     a					=inputs.pop('''labels'''						)
     a		,				a					=self._compute_loss(__A			,   __A			,   __A						)
     return loss


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A = None			,   ) ->					Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
     a					=self._prepare_inputs(__A						)

     a					={
         '''max_length''': self.data_args.val_max_target_length
         if self.data_args is not None
         else self.config.max_length,
         '''num_beams''': self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
     }

     if self.args.predict_with_generate and not self.args.prediction_loss_only:
         a					=self.model.generate(
             inputs['''input_ids''']			,   attention_mask=inputs['''attention_mask''']			,   **__A			,   )
         # in case the batch is shorter than max length, the output should be padded
         if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
             a					=self._pad_tensors_to_max_len(__A			,   gen_kwargs['''max_length''']						)

     a					=inputs.pop('''labels'''						)
     with torch.no_grad():
         # compute loss on predict data
         a		,				a					=self._compute_loss(__A			,   __A			,   __A						)

     a					=loss.mean().detach()
     if self.args.prediction_loss_only:
         return (loss, None, None)

     a					=generated_tokens if self.args.predict_with_generate else logits

     if labels.shape[-1] < gen_kwargs["max_length"]:
         a					=self._pad_tensors_to_max_len(__A			,   gen_kwargs['''max_length''']						)

     return (loss, logits, labels)





 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A						) ->					List[Any]:
     # If PAD token is not defined at least EOS token has to be defined
     a					=self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

     if pad_token_id is None:
         raise ValueError(
             '''Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'''
             f''' padded to `max_length`={max_length}'''						)

     a					=pad_token_id * torch.ones(
         (tensor.shape[0], max_length)			,   dtype=tensor.dtype			,   device=tensor.device						)
     a					=tensor
     return padded_tensor



"""simple docstring"""







import inspect
import unittest

from transformers import YolosConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
       import torch
       from torch import nn

       from transformers import YolosForObjectDetection, YolosModel
       from transformers.models.yolos.modeling_yolos import YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor






class 				__A    :






 """simple docstring"""




 def __init__(		self			,   __A			,   __A=13			,   __A=[30, 30]			,   __A=2			,   __A=3			,   __A=True			,   __A=True			,   __A=32			,   __A=5			,   __A=4			,   __A=37			,   __A="gelu"			,   __A=0.1			,   __A=0.1			,   __A=10			,   __A=0.02			,   __A=3			,   __A=None			,   __A=8			,   __A=10			,   ) ->					List[Any]:
     a					=parent
     a					=batch_size
     a					=image_size
     a					=patch_size
     a					=num_channels
     a					=is_training
     a					=use_labels
     a					=hidden_size
     a					=num_hidden_layers
     a					=num_attention_heads
     a					=intermediate_size
     a					=hidden_act
     a					=hidden_dropout_prob
     a					=attention_probs_dropout_prob
     a					=type_sequence_label_size
     a					=initializer_range
     a					=num_labels
     a					=scope
     a					=n_targets
     a					=num_detection_tokens
     # we set the expected sequence length (which is used in several tests)
     # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token) + num_detection_tokens
     a					=(image_size[1] // patch_size) * (image_size[0] // patch_size)
     a					=num_patches + 1 + self.num_detection_tokens


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
     a					=floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]						)

     a					=None
     if self.use_labels:
         # labels is a list of Dict (each Dict being the labels for a given example in the batch)
         a					=[]
         for i in range(self.batch_size						):
             a					={}
             a					=torch.randint(
                 high=self.num_labels			,   size=(self.n_targets,)			,   device=__A						)
             a					=torch.rand(self.n_targets			,   4			,   device=__A						)
             labels.append(__A						)

     a					=self.get_config()

     return config, pixel_values, labels


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					int:
     return YolosConfig(
         image_size=self.image_size			,   patch_size=self.patch_size			,   num_channels=self.num_channels			,   hidden_size=self.hidden_size			,   num_hidden_layers=self.num_hidden_layers			,   num_attention_heads=self.num_attention_heads			,   intermediate_size=self.intermediate_size			,   hidden_act=self.hidden_act			,   hidden_dropout_prob=self.hidden_dropout_prob			,   attention_probs_dropout_prob=self.attention_probs_dropout_prob			,   is_decoder=__A			,   initializer_range=self.initializer_range			,   num_detection_tokens=self.num_detection_tokens			,   num_labels=self.num_labels			,   )


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A						) ->					List[Any]:
     a					=YolosModel(config=__A						)
     model.to(__A						)
     model.eval()
     a					=model(__A						)
     self.parent.assertEqual(
         result.last_hidden_state.shape			,   (self.batch_size, self.expected_seq_len, self.hidden_size)						)


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A						) ->					Dict:
     a					=YolosForObjectDetection(__A						)
     model.to(__A						)
     model.eval()

     a					=model(pixel_values=__A						)
     a					=model(__A						)

     self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_detection_tokens, self.num_labels + 1)						)
     self.parent.assertEqual(result.pred_boxes.shape			,   (self.batch_size, self.num_detection_tokens, 4)						)

     a					=model(pixel_values=__A			,   labels=__A						)

     self.parent.assertEqual(result.loss.shape			,   ()						)
     self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_detection_tokens, self.num_labels + 1)						)
     self.parent.assertEqual(result.pred_boxes.shape			,   (self.batch_size, self.num_detection_tokens, 4)						)





 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Union[str, Any]:
     a					=self.prepare_config_and_inputs()
     a		,				a		,				a					=config_and_inputs
     a					={'''pixel_values''': pixel_values}
     return config, inputs_dict






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






 """simple docstring"""




 __lowerCAmelCase			 =  (YolosModel, YolosForObjectDetection) if is_torch_available() else ()
 __lowerCAmelCase			 =  (
     {"feature-extraction": YolosModel, "object-detection": YolosForObjectDetection} if is_torch_available() else {}
 )

 __lowerCAmelCase			 =  False
 __lowerCAmelCase			 =  False
 __lowerCAmelCase			 =  False
 __lowerCAmelCase			 =  False


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A=False						) ->					Any:
     a					=super()._prepare_for_class(__A			,   __A			,   return_labels=__A						)

     if return_labels:
         if model_class.__name__ == "YolosForObjectDetection":
             a					=[]
             for i in range(self.model_tester.batch_size						):
                 a					={}
                 a					=torch.ones(
                     size=(self.model_tester.n_targets,)			,   device=__A			,   dtype=torch.long						)
                 a					=torch.ones(
                     self.model_tester.n_targets			,   4			,   device=__A			,   dtype=torch.float						)
                 labels.append(__A						)
             a					=labels

     return inputs_dict


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[str]:
     a					=YolosModelTester(self						)
     a					=ConfigTester(self			,   config_class=__A			,   has_text_modality=__A			,   hidden_size=37						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Dict:
     self.config_tester.run_common_tests()


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     # YOLOS does not use inputs_embeds
     pass


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
     a		,				a					=self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         a					=model_class(__A						)
         self.assertIsInstance(model.get_input_embeddings()			,   (nn.Module)						)
         a					=model.get_output_embeddings()
         self.assertTrue(x is None or isinstance(__A			,   nn.Linear						)						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[int]:
     a		,				a					=self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         a					=model_class(__A						)
         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]			,   __A						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
     a					=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__A						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     a		,				a					=self.model_tester.prepare_config_and_inputs_for_common()
     a					=True

     # in YOLOS, the seq_len is different
     a					=self.model_tester.expected_seq_len
     for model_class in self.all_model_classes:
         a					=True
         a					=False
         a					=True
         a					=model_class(__A						)
         model.to(__A						)
         model.eval()
         with torch.no_grad():
             a					=model(**self._prepare_for_class(__A			,   __A						)						)
         a					=outputs.attentions
         self.assertEqual(len(__A						)			,   self.model_tester.num_hidden_layers						)

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

         self.assertListEqual(
             list(attentions[0].shape[-3:]						)			,   [self.model_tester.num_attention_heads, seq_len, seq_len]			,   )
         a					=len(__A						)

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

         a					=1
         self.assertEqual(out_len + added_hidden_states			,   len(__A						)						)

         a					=outputs.attentions

         self.assertEqual(len(__A						)			,   self.model_tester.num_hidden_layers						)
         self.assertListEqual(
             list(self_attentions[0].shape[-3:]						)			,   [self.model_tester.num_attention_heads, seq_len, seq_len]			,   )


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					str:
     def check_hidden_states_output(__A			,   __A			,   __A						):
         a					=model_class(__A						)
         model.to(__A						)
         model.eval()

         with torch.no_grad():
             a					=model(**self._prepare_for_class(__A			,   __A						)						)

         a					=outputs.hidden_states

         a					=getattr(
             self.model_tester			,   '''expected_num_hidden_layers'''			,   self.model_tester.num_hidden_layers + 1						)
         self.assertEqual(len(__A						)			,   __A						)

         # YOLOS has a different seq_length
         a					=self.model_tester.expected_seq_len

         self.assertListEqual(
             list(hidden_states[0].shape[-2:]						)			,   [seq_length, self.model_tester.hidden_size]			,   )

     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(__A			,   __A			,   __A						)

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

         check_hidden_states_output(__A			,   __A			,   __A						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					int:
     a					=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_object_detection(*__A						)





 @slow
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         a					=YolosModel.from_pretrained(__A						)
         self.assertIsNotNone(__A						)

def _A						(						):







    """simple docstring"""
    a					=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
    return image






@require_torch
@require_vision
class 				__A    (	unittest.TestCase      ):






 """simple docstring"""




 @cached_property
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Dict:
     return AutoImageProcessor.from_pretrained('''hustvl/yolos-small'''						) if is_vision_available() else None





 @slow
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Union[str, Any]:
     a					=YolosForObjectDetection.from_pretrained('''hustvl/yolos-small'''						).to(__A						)

     a					=self.default_image_processor
     a					=prepare_img()
     a					=image_processor(images=__A			,   return_tensors='''pt'''						).to(__A						)

     # forward pass
     with torch.no_grad():
         a					=model(inputs.pixel_values						)

     # verify outputs
     a					=torch.Size((1, 100, 92)						)
     self.assertEqual(outputs.logits.shape			,   __A						)

     a					=torch.tensor(
         [[-24.0_248, -10.3_024, -14.8_290], [-42.0_392, -16.8_200, -27.4_334], [-27.2_743, -11.8_154, -18.7_148]]			,   device=__A			,   )
     a					=torch.tensor(
         [[0.2_559, 0.5_455, 0.4_706], [0.2_989, 0.7_279, 0.1_875], [0.7_732, 0.4_017, 0.4_462]]			,   device=__A						)
     self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]			,   __A			,   atol=1E-4						)						)
     self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3]			,   __A			,   atol=1E-4						)						)

     # verify postprocessing
     a					=image_processor.post_process_object_detection(
         __A			,   threshold=0.3			,   target_sizes=[image.size[::-1]]						)[0]
     a					=torch.tensor([0.9_994, 0.9_790, 0.9_964, 0.9_972, 0.9_861]						).to(__A						)
     a					=[75, 75, 17, 63, 17]
     a					=torch.tensor([335.0_609, 79.3_848, 375.4_216, 187.2_495]						).to(__A						)

     self.assertEqual(len(results['''scores''']						)			,   5						)
     self.assertTrue(torch.allclose(results['''scores''']			,   __A			,   atol=1E-4						)						)
     self.assertSequenceEqual(results['''labels'''].tolist()			,   __A						)
     self.assertTrue(torch.allclose(results['''boxes'''][0, :]			,   __A						)						)