




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 		lowercase				:

     def __init__(    self						,	_a						,	_a=13						,	_a=10						,	_a=3						,	_a=2						,	_a=2						,	_a=2						,	_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=0.9						,	_a=None						,	)		->				Union[str, Any]:
           _A				:	Optional[int]     =		parent
           _A				:	int     =		batch_size
           _A				:	Optional[int]     =		image_size
           _A				:	Optional[int]     =		num_channels
           _A				:	Tuple     =		patch_size
           _A				:	Tuple     =		tubelet_size
           _A				:	Any     =		num_frames
           _A				:	List[Any]     =		is_training
           _A				:	int     =		use_labels
           _A				:	str     =		hidden_size
           _A				:	Tuple     =		num_hidden_layers
           _A				:	Dict     =		num_attention_heads
           _A				:	Optional[int]     =		intermediate_size
           _A				:	Any     =		hidden_act
           _A				:	int     =		hidden_dropout_prob
           _A				:	Optional[Any]     =		attention_probs_dropout_prob
           _A				:	Optional[Any]     =		type_sequence_label_size
           _A				:	Union[str, Any]     =		initializer_range
           _A				:	Dict     =		mask_ratio
           _A				:	List[str]     =		scope

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

           # use this variable to define bool_masked_pos
           _A				:	Optional[Any]     =		int(mask_ratio * self.seq_length							)

     def    a__      (    self							)		->				Optional[int]:
           _A				:	Optional[Any]     =		floats_tensor(
               [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]							)

           _A				:	List[Any]     =		None
           if self.use_labels:
                 _A				:	Any     =		ids_tensor([self.batch_size]						,	self.type_sequence_label_size							)

           _A				:	Any     =		self.get_config()

           return config, pixel_values, labels

     def    a__      (    self							)		->				Union[str, Any]:
           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						,	_a						,	_a						,	_a							)		->				Any:
           _A				:	Dict     =		VideoMAEModel(config=_a							)
           model.to(_a							)
           model.eval()
           _A				:	Optional[int]     =		model(_a							)
           self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.hidden_size)							)

     def    a__      (    self						,	_a						,	_a						,	_a							)		->				Any:
           _A				:	Optional[Any]     =		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
           _A				:	Tuple     =		torch.ones((self.num_masks,)							)
           _A				:	int     =		torch.cat([mask, torch.zeros(self.seq_length - mask.size(0							)							)]							)
           _A				:	Optional[Any]     =		mask.expand(self.batch_size						,	-1							).bool()

           _A				:	Any     =		model(_a						,	_a							)
           # model only returns predictions for masked patches
           _A				:	int     =		mask.sum().item()
           _A				:	Tuple     =		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							)		->				List[str]:
           _A				:	Optional[Any]     =		self.prepare_config_and_inputs()
           _A				:	Optional[int]     =		config_and_inputs
           _A				:	int     =		{"""pixel_values""": pixel_values}
           return config, inputs_dict







@require_torch
class 		lowercase				(  _UpperCamelCase,_UpperCamelCase,unittest.TestCase						):
     _a									=		(
         (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
     )
     _a									=		(
         {'feature-extraction': VideoMAEModel, 'video-classification': VideoMAEForVideoClassification}
         if is_torch_available()
         else {}
     )

     _a									=		False
     _a									=		False
     _a									=		False
     _a									=		False

     def    a__      (    self							)		->				Tuple:
           _A				:	Any     =		VideoMAEModelTester(self							)
           _A				:	List[str]     =		ConfigTester(self						,	config_class=_a						,	has_text_modality=_a						,	hidden_size=37							)

     def    a__      (    self						,	_a						,	_a						,	_a=False							)		->				Union[str, Any]:
           _A				:	Any     =		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
                 _A				:	Any     =		torch.ones((self.model_tester.num_masks,)							)
                 _A				:	str     =		torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0							)							)]							)
                 _A				:	Dict     =		mask.expand(self.model_tester.batch_size						,	-1							).bool()
                 _A				:	int     =		bool_masked_pos.to(_a							)

           if return_labels:
                 if model_class in [
                     *get_values(_a							),
                 ]:
                       _A				:	List[Any]     =		torch.zeros(
                           self.model_tester.batch_size						,	dtype=torch.long						,	device=_a							)

           return inputs_dict

     def    a__      (    self							)		->				str:
           self.config_tester.run_common_tests()

     @unittest.skip(reason="""VideoMAE does not use inputs_embeds"""							)
     def    a__      (    self							)		->				Any:
           pass

     def    a__      (    self							)		->				Union[str, Any]:
           _A				:	str     =		self.model_tester.prepare_config_and_inputs_for_common()

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

     def    a__      (    self							)		->				Dict:
           _A				:	Optional[int]     =		self.model_tester.prepare_config_and_inputs_for_common()

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

                 _A				:	Union[str, Any]     =		["""pixel_values"""]
                 self.assertListEqual(arg_names[:1]						,	_a							)

     def    a__      (    self							)		->				Optional[Any]:
           _A				:	Tuple     =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*_a							)

     def    a__      (    self							)		->				str:
           _A				:	Tuple     =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*_a							)

     @slow
     def    a__      (    self							)		->				Dict:







           for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _A				:	Dict     =		VideoMAEModel.from_pretrained(_a							)
                 self.assertIsNotNone(_a							)

     def    a__      (    self							)		->				int:







           if not self.has_attentions:
                 pass

           else:
                 _A				:	int     =		self.model_tester.prepare_config_and_inputs_for_common()
                 _A				:	Optional[Any]     =		True

                 for model_class in self.all_model_classes:
                       _A				:	List[str]     =		self.model_tester.seq_length - self.model_tester.num_masks
                       _A				:	Tuple     =		(
                           num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
                       )

                       _A				:	Optional[Any]     =		True
                       _A				:	str     =		False
                       _A				:	Optional[int]     =		True
                       _A				:	int     =		model_class(_a							)
                       model.to(_a							)
                       model.eval()
                       with torch.no_grad():
                             _A				:	Tuple     =		model(**self._prepare_for_class(_a						,	_a							)							)
                       _A				:	str     =		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				:	List[str]     =		True
                       _A				:	str     =		model_class(_a							)
                       model.to(_a							)
                       model.eval()
                       with torch.no_grad():
                             _A				:	Optional[Any]     =		model(**self._prepare_for_class(_a						,	_a							)							)
                       _A				:	Tuple     =		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				:	Tuple     =		len(_a							)

                       # Check attention is always last and order is fine
                       _A				:	Dict     =		True
                       _A				:	Tuple     =		True
                       _A				:	str     =		model_class(_a							)
                       model.to(_a							)
                       model.eval()
                       with torch.no_grad():
                             _A				:	Optional[int]     =		model(**self._prepare_for_class(_a						,	_a							)							)

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

                       _A				:	Optional[Any]     =		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							)		->				List[str]:







           def check_hidden_states_output(_a						,	_a						,	_a							):
                 _A				:	int     =		model_class(_a							)
                 model.to(_a							)
                 model.eval()

                 with torch.no_grad():
                       _A				:	str     =		model(**self._prepare_for_class(_a						,	_a							)							)

                 _A				:	Optional[int]     =		outputs.hidden_states
                 _A				:	Union[str, Any]     =		self.model_tester.num_hidden_layers + 1
                 self.assertEqual(len(_a							)						,	_a							)

                 _A				:	Any     =		self.model_tester.seq_length - self.model_tester.num_masks
                 _A				:	Dict     =		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]						,	)

           _A				:	Optional[int]     =		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 _A				:	Union[str, Any]     =		True
                 check_hidden_states_output(_a						,	_a						,	_a							)

                 # check that output_hidden_states also work using config
                 del inputs_dict["output_hidden_states"]
                 _A				:	Union[str, Any]     =		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							)		->				Optional[Any]:
           pass







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







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

     @cached_property
     def    a__      (    self							)		->				List[str]:
           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							)		->				Optional[Any]:
           _A				:	Union[str, Any]     =		VideoMAEForVideoClassification.from_pretrained("""MCG-NJU/videomae-base-finetuned-kinetics"""							).to(
               _a							)

           _A				:	Any     =		self.default_image_processor
           _A				:	str     =		prepare_video()
           _A				:	Dict     =		image_processor(_a						,	return_tensors="""pt"""							).to(_a							)

           # forward pass
           with torch.no_grad():
                 _A				:	Optional[int]     =		model(**_a							)

           # verify the logits
           _A				:	Optional[int]     =		torch.Size((1, 400)							)
           self.assertEqual(outputs.logits.shape						,	_a							)

           _A				:	Any     =		torch.tensor([0.3669, -0.0688, -0.2421]							).to(_a							)

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



     @slow
     def    a__      (    self							)		->				Any:
           _A				:	Tuple     =		VideoMAEForPreTraining.from_pretrained("""MCG-NJU/videomae-base-short"""							).to(_a							)

           _A				:	int     =		self.default_image_processor
           _A				:	Optional[int]     =		prepare_video()
           _A				:	int     =		image_processor(_a						,	return_tensors="""pt"""							).to(_a							)

           # add boolean mask, indicating which patches to mask
           _A				:	Tuple     =		hf_hub_download(repo_id="""hf-internal-testing/bool-masked-pos"""						,	filename="""bool_masked_pos.pt"""							)
           _A				:	int     =		torch.load(_a							)

           # forward pass
           with torch.no_grad():
                 _A				:	int     =		model(**_a							)

           # verify the logits
           _A				:	str     =		torch.Size([1, 1408, 1536]							)
           _A				:	Union[str, Any]     =		torch.tensor(
               [[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]						,	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`)
           _A				:	Tuple     =		torch.tensor([0.5142]						,	device=_a							)
           self.assertTrue(torch.allclose(outputs.loss						,	_a						,	atol=1e-4							)							)

           # verify the loss (`config.norm_pix_loss` = `False`)
           _A				:	Optional[Any]     =		VideoMAEForPreTraining.from_pretrained("""MCG-NJU/videomae-base-short"""						,	norm_pix_loss=_a							).to(
               _a							)

           with torch.no_grad():
                 _A				:	Union[str, Any]     =		model(**_a							)

           _A				:	Tuple     =		torch.tensor(torch.tensor([0.6469]							)						,	device=_a							)
           self.assertTrue(torch.allclose(outputs.loss						,	_a						,	atol=1e-4							)							)







# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from accelerate import PartialState
from accelerate.utils.operations import broadcast, gather, gather_object, pad_across_processes, reduce





def 				A__ (snake_case   :		List[Any]				)     ->						Optional[Any]:
					return (torch.arange(state.num_processes				) + 1.0 + (state.num_processes * state.process_index)).to(state.device				)





def 				A__ (snake_case   :		int				)     ->						Any:
					__UpperCamelCase					:	List[str]             =      create_tensor(snake_case				)
					__UpperCamelCase					:	List[Any]             =      gather(snake_case				)
					assert gathered_tensor.tolist() == list(range(1		,	state.num_processes**2 + 1				)				)





def 				A__ (snake_case   :		Optional[Any]				)     ->						Dict:
					__UpperCamelCase					:	Any             =      [state.process_index]
					__UpperCamelCase					:	int             =      gather_object(snake_case				)
					assert len(snake_case				) == state.num_processes, F'''{gathered_obj}, {len(snake_case				)} != {state.num_processes}'''
					assert gathered_obj == list(range(state.num_processes				)				), F'''{gathered_obj} != {list(range(state.num_processes				)				)}'''





def 				A__ (snake_case   :		Any				)     ->						Optional[Any]:
					__UpperCamelCase					:	Optional[Any]             =      create_tensor(snake_case				)
					__UpperCamelCase					:	Any             =      broadcast(snake_case				)
					assert broadcasted_tensor.shape == torch.Size([state.num_processes]				)
					assert broadcasted_tensor.tolist() == list(range(1		,	state.num_processes + 1				)				)





def 				A__ (snake_case   :		Tuple				)     ->						Dict:
					# We need to pad the tensor with one more element if we are the main process
					# to ensure that we can pad
					if state.is_main_process:
										__UpperCamelCase					:	List[Any]             =      torch.arange(state.num_processes + 1				).to(state.device				)
					else:
										__UpperCamelCase					:	List[str]             =      torch.arange(state.num_processes				).to(state.device				)
					__UpperCamelCase					:	List[str]             =      pad_across_processes(snake_case				)
					assert padded_tensor.shape == torch.Size([state.num_processes + 1]				)
					if not state.is_main_process:
										assert padded_tensor.tolist() == list(range(0		,	state.num_processes				)				) + [0]





def 				A__ (snake_case   :		Tuple				)     ->						Dict:
					# For now runs on only two processes
					if state.num_processes != 2:
										return
					__UpperCamelCase					:	Optional[int]             =      create_tensor(snake_case				)
					__UpperCamelCase					:	str             =      reduce(snake_case		,	"""sum"""				)
					__UpperCamelCase					:	Optional[int]             =      torch.tensor([4.0, 6]				).to(state.device				)
					assert torch.allclose(snake_case		,	snake_case				), F'''{reduced_tensor} != {truth_tensor}'''





def 				A__ (snake_case   :		Optional[Any]				)     ->						Optional[Any]:
					# For now runs on only two processes
					if state.num_processes != 2:
										return
					__UpperCamelCase					:	List[Any]             =      create_tensor(snake_case				)
					__UpperCamelCase					:	str             =      reduce(snake_case		,	"""mean"""				)
					__UpperCamelCase					:	Union[str, Any]             =      torch.tensor([2.0, 3]				).to(state.device				)
					assert torch.allclose(snake_case		,	snake_case				), F'''{reduced_tensor} != {truth_tensor}'''





def 				A__ (snake_case   :		Tuple				)     ->						int:
					# For xla_spawn (TPUs)
					main()





def 				A__ ()     ->						List[Any]:
					__UpperCamelCase					:	List[Any]             =      PartialState()
					state.print(F'''State: {state}'''				)
					state.print("""testing gather"""				)
					test_gather(snake_case				)
					state.print("""testing gather_object"""				)
					test_gather_object(snake_case				)
					state.print("""testing broadcast"""				)
					test_broadcast(snake_case				)
					state.print("""testing pad_across_processes"""				)
					test_pad_across_processes(snake_case				)
					state.print("""testing reduce_sum"""				)
					test_reduce_sum(snake_case				)
					state.print("""testing reduce_mean"""				)
					test_reduce_mean(snake_case				)


if __name__ == "__main__":
	main()






