





import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
from transformers.models.auto import get_values
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, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import (
        MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
        VideoMAEForPreTraining,
        VideoMAEForVideoClassification,
        VideoMAEModel,
    )
    from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from transformers import VideoMAEImageProcessor

class 		SCREAMING_SNAKE_CASE__    :


      def __init__(self	:				Any    ,      a__	:				str    ,      a__	:				Tuple=13    ,      a__	:				int=10    ,      a__	:				Tuple=3    ,      a__	:				Tuple=2    ,      a__	:				List[str]=2    ,      a__	:				Union[str, Any]=2    ,      a__	:				Dict=True    ,      a__	:				Union[str, Any]=True    ,      a__	:				List[str]=32    ,      a__	:				Dict=5    ,      a__	:				Tuple=4    ,      a__	:				List[str]=37    ,      a__	:				Dict="gelu"    ,      a__	:				Dict=0.1    ,      a__	:				int=0.1    ,      a__	:				Optional[Any]=10    ,      a__	:				List[str]=0.0_2    ,      a__	:				str=0.9    ,      a__	:				Optional[Any]=None    ,      ):





         """simple docstring"""



         __snake_case								=       parent
         __snake_case								=       batch_size
         __snake_case								=       image_size
         __snake_case								=       num_channels
         __snake_case								=       patch_size
         __snake_case								=       tubelet_size
         __snake_case								=       num_frames
         __snake_case								=       is_training
         __snake_case								=       use_labels
         __snake_case								=       hidden_size
         __snake_case								=       num_hidden_layers
         __snake_case								=       num_attention_heads
         __snake_case								=       intermediate_size
         __snake_case								=       hidden_act
         __snake_case								=       hidden_dropout_prob
         __snake_case								=       attention_probs_dropout_prob
         __snake_case								=       type_sequence_label_size
         __snake_case								=       initializer_range
         __snake_case								=       mask_ratio
         __snake_case								=       scope

         # in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
         __snake_case								=       (image_size // patch_size) ** 2
         __snake_case								=       (num_frames // tubelet_size) * self.num_patches_per_frame

         # use this variable to define bool_masked_pos
         __snake_case								=       int(mask_ratio * self.seq_length				)


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





         """simple docstring"""



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

         __snake_case								=       None
         if self.use_labels:
            __snake_case								=       ids_tensor([self.batch_size]    ,      self.type_sequence_label_size				)

         __snake_case								=       self.get_config()

         return config, pixel_values, labels


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         return VideoMAEConfig(
             image_size=self.image_size    ,      patch_size=self.patch_size    ,      num_channels=self.num_channels    ,      num_frames=self.num_frames    ,      tubelet_size=self.tubelet_size    ,      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    ,      )


      def        a			(self	:				List[Any]    ,      a__	:				int    ,      a__	:				Tuple    ,      a__	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       VideoMAEModel(config=a__				)
         model.to(a__				)
         model.eval()
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.hidden_size)				)


      def        a			(self	:				str    ,      a__	:				Union[str, Any]    ,      a__	:				Optional[Any]    ,      a__	:				Tuple				):





         """simple docstring"""



         __snake_case								=       VideoMAEForPreTraining(a__				)
         model.to(a__				)
         model.eval()
         # important: each video needs to have the same number of masked patches
         # hence we define a single mask, which we then repeat for each example in the batch
         __snake_case								=       torch.ones((self.num_masks,)				)
         __snake_case								=       torch.cat([mask, torch.zeros(self.seq_length - mask.size(0				)				)]				)
         __snake_case								=       mask.expand(self.batch_size    ,      -1				).bool()

         __snake_case								=       model(a__    ,      a__				)
         # model only returns predictions for masked patches
         __snake_case								=       mask.sum().item()
         __snake_case								=       3 * self.tubelet_size * self.patch_size**2
         self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, num_masked_patches, decoder_num_labels)				)


      def        a			(self	:				Any				):





         """simple docstring"""



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






@require_torch
class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    Dict							=   (
          (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
      )
      A_						:    Any							=   (
          {'feature-extraction': VideoMAEModel, 'video-classification': VideoMAEForVideoClassification}
          if is_torch_available()
          else {}
      )

      A_						:    str							=   False
      A_						:    Tuple							=   False
      A_						:    List[str]							=   False
      A_						:    Tuple							=   False


      def        a			(self	:				List[str]				):





         """simple docstring"""



         __snake_case								=       VideoMAEModelTester(self				)
         __snake_case								=       ConfigTester(self    ,      config_class=a__    ,      has_text_modality=a__    ,      hidden_size=37				)


      def        a			(self	:				Tuple    ,      a__	:				Any    ,      a__	:				Tuple    ,      a__	:				Optional[Any]=False				):





         """simple docstring"""



         __snake_case								=       copy.deepcopy(a__				)

         if model_class == VideoMAEForPreTraining:
            # important: each video needs to have the same number of masked patches
            # hence we define a single mask, which we then repeat for each example in the batch
            __snake_case								=       torch.ones((self.model_tester.num_masks,)				)
            __snake_case								=       torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0				)				)]				)
            __snake_case								=       mask.expand(self.model_tester.batch_size    ,      -1				).bool()
            __snake_case								=       bool_masked_pos.to(a__				)

         if return_labels:
            if model_class in [
                *get_values(a__				),
            ]:
               __snake_case								=       torch.zeros(
                   self.model_tester.batch_size    ,      dtype=torch.long    ,      device=a__				)

         return inputs_dict


      def        a			(self	:				Dict				):





         """simple docstring"""



         self.config_tester.run_common_tests()


      @unittest.skip(reason='''VideoMAE does not use inputs_embeds'''				)
      def        a			(self	:				Tuple				):





         """simple docstring"""



         pass


      def        a			(self	:				int				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()

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


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()

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

            __snake_case								=       ['''pixel_values''']
            self.assertListEqual(arg_names[:1]    ,      a__				)


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__				)


      def        a			(self	:				Any				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*a__				)


      @slow
      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __snake_case								=       VideoMAEModel.from_pretrained(a__				)
            self.assertIsNotNone(a__				)


      def        a			(self	:				List[str]				):





         """simple docstring"""



         if not self.has_attentions:
            pass

         else:
            __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()
            __snake_case								=       True

            for model_class in self.all_model_classes:
               __snake_case								=       self.model_tester.seq_length - self.model_tester.num_masks
               __snake_case								=       (
                   num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
               )

               __snake_case								=       True
               __snake_case								=       False
               __snake_case								=       True
               __snake_case								=       model_class(a__				)
               model.to(a__				)
               model.eval()
               with torch.no_grad():
                  __snake_case								=       model(**self._prepare_for_class(a__    ,      a__				)				)
               __snake_case								=       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"]
               __snake_case								=       True
               __snake_case								=       model_class(a__				)
               model.to(a__				)
               model.eval()
               with torch.no_grad():
                  __snake_case								=       model(**self._prepare_for_class(a__    ,      a__				)				)
               __snake_case								=       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]    ,      )
               __snake_case								=       len(a__				)

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

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

               __snake_case								=       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        a			(self	:				Any				):





         """simple docstring"""



         def check_hidden_states_output(a__	:				Optional[int]    ,      a__	:				str    ,      a__	:				int				):
            __snake_case								=       model_class(a__				)
            model.to(a__				)
            model.eval()

            with torch.no_grad():
               __snake_case								=       model(**self._prepare_for_class(a__    ,      a__				)				)

            __snake_case								=       outputs.hidden_states
            __snake_case								=       self.model_tester.num_hidden_layers + 1
            self.assertEqual(len(a__				)    ,      a__				)

            __snake_case								=       self.model_tester.seq_length - self.model_tester.num_masks
            __snake_case								=       num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

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

         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            __snake_case								=       True
            check_hidden_states_output(a__    ,      a__    ,      a__				)

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

            check_hidden_states_output(a__    ,      a__    ,      a__				)


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





         """simple docstring"""



         pass






def lowerCamelCase__   (							)				->    Union[str, Any]:
   __snake_case								=       hf_hub_download(
       repo_id='''hf-internal-testing/spaghetti-video'''	,				filename='''eating_spaghetti.npy'''	,				repo_type='''dataset'''	)
   __snake_case								=       np.load(snake_case_	)
   return list(snake_case_	)







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


      @cached_property
      def        a			(self	:				int				):





         """simple docstring"""



         return (
             VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]    ,      image_std=[0.5, 0.5, 0.5]				)
             if is_vision_available()
             else None
         )


      @slow
      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         __snake_case								=       VideoMAEForVideoClassification.from_pretrained('''MCG-NJU/videomae-base-finetuned-kinetics'''				).to(
             a__				)

         __snake_case								=       self.default_image_processor
         __snake_case								=       prepare_video()
         __snake_case								=       image_processor(a__    ,      return_tensors='''pt'''				).to(a__				)

         # forward pass
         with torch.no_grad():
            __snake_case								=       model(**a__				)

         # verify the logits
         __snake_case								=       torch.Size((1, 400)				)
         self.assertEqual(outputs.logits.shape    ,      a__				)

         __snake_case								=       torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]				).to(a__				)

         self.assertTrue(torch.allclose(outputs.logits[0, :3]    ,      a__    ,      atol=1E-4				)				)


      @slow
      def        a			(self	:				Tuple				):





         """simple docstring"""



         __snake_case								=       VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short'''				).to(a__				)

         __snake_case								=       self.default_image_processor
         __snake_case								=       prepare_video()
         __snake_case								=       image_processor(a__    ,      return_tensors='''pt'''				).to(a__				)

         # add boolean mask, indicating which patches to mask
         __snake_case								=       hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''    ,      filename='''bool_masked_pos.pt'''				)
         __snake_case								=       torch.load(a__				)

         # forward pass
         with torch.no_grad():
            __snake_case								=       model(**a__				)

         # verify the logits
         __snake_case								=       torch.Size([1, 1408, 1536]				)
         __snake_case								=       torch.tensor(
             [[0.7_9_9_4, 0.9_6_1_2, 0.8_5_0_8], [0.7_4_0_1, 0.8_9_5_8, 0.8_3_0_2], [0.5_8_6_2, 0.7_4_6_8, 0.7_3_2_5]]    ,      device=a__				)
         self.assertEqual(outputs.logits.shape    ,      a__				)
         self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]    ,      a__    ,      atol=1E-4				)				)

         # verify the loss (`config.norm_pix_loss` = `True`)
         __snake_case								=       torch.tensor([0.5_1_4_2]    ,      device=a__				)
         self.assertTrue(torch.allclose(outputs.loss    ,      a__    ,      atol=1E-4				)				)

         # verify the loss (`config.norm_pix_loss` = `False`)
         __snake_case								=       VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short'''    ,      norm_pix_loss=a__				).to(
             a__				)

         with torch.no_grad():
            __snake_case								=       model(**a__				)

         __snake_case								=       torch.tensor(torch.tensor([0.6_4_6_9]				)    ,      device=a__				)
         self.assertTrue(torch.allclose(outputs.loss    ,      a__    ,      atol=1E-4				)				)








class 		SCREAMING_SNAKE_CASE__    :


      def __init__(self	:				str    ,      a__	:				list				):





         """simple docstring"""



         __snake_case								=       set_counts
         __snake_case								=       max(a__				)
         __snake_case								=       len(a__				)
         __snake_case								=       [1] * num_sets
         __snake_case								=       list(range(a__				)				)


      def        a			(self	:				str    ,      a__	:				int    ,      a__	:				int				):





         """simple docstring"""



         __snake_case								=       self.get_parent(a__				)
         __snake_case								=       self.get_parent(a__				)

         if src_parent == dst_parent:
            return False

         if self.ranks[dst_parent] >= self.ranks[src_parent]:
            self.set_counts[dst_parent] += self.set_counts[src_parent]
            __snake_case								=       0
            __snake_case								=       dst_parent
            if self.ranks[dst_parent] == self.ranks[src_parent]:
               self.ranks[dst_parent] += 1
            __snake_case								=       self.set_counts[dst_parent]
         else:
            self.set_counts[src_parent] += self.set_counts[dst_parent]
            __snake_case								=       0
            __snake_case								=       src_parent
            __snake_case								=       self.set_counts[src_parent]

         __snake_case								=       max(self.max_set    ,      a__				)
         return True


      def        a			(self	:				Union[str, Any]    ,      a__	:				int				):





         """simple docstring"""



         if self.parents[disj_set] == disj_set:
            return disj_set
         __snake_case								=       self.get_parent(self.parents[disj_set]				)
         return self.parents[disj_set]

