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







							def __init__(				self     :    Optional[Any]	,					__snake_case     :    Tuple	,					__snake_case     :    Union[str, Any]=13	,					__snake_case     :    int=10	,					__snake_case     :    Optional[Any]=3	,					__snake_case     :    Union[str, Any]=2	,					__snake_case     :    Tuple=2	,					__snake_case     :    Union[str, Any]=2	,					__snake_case     :    int=True	,					__snake_case     :    Dict=True	,					__snake_case     :    Optional[int]=32	,					__snake_case     :    int=5	,					__snake_case     :    List[Any]=4	,					__snake_case     :    Union[str, Any]=37	,					__snake_case     :    List[str]="gelu"	,					__snake_case     :    int=0.1	,					__snake_case     :    Optional[int]=0.1	,					__snake_case     :    str=10	,					__snake_case     :    Dict=0.0_2	,					__snake_case     :    Union[str, Any]=0.9	,					__snake_case     :    Any=None	,					):
										lowerCamelCase						:Dict						  =     parent
										lowerCamelCase						:List[Any]						  =     batch_size
										lowerCamelCase						:str						  =     image_size
										lowerCamelCase						:int						  =     num_channels
										lowerCamelCase						:Any						  =     patch_size
										lowerCamelCase						:Tuple						  =     tubelet_size
										lowerCamelCase						:int						  =     num_frames
										lowerCamelCase						:Any						  =     is_training
										lowerCamelCase						:List[str]						  =     use_labels
										lowerCamelCase						:Dict						  =     hidden_size
										lowerCamelCase						:int						  =     num_hidden_layers
										lowerCamelCase						:List[str]						  =     num_attention_heads
										lowerCamelCase						:Optional[int]						  =     intermediate_size
										lowerCamelCase						:Union[str, Any]						  =     hidden_act
										lowerCamelCase						:Dict						  =     hidden_dropout_prob
										lowerCamelCase						:Tuple						  =     attention_probs_dropout_prob
										lowerCamelCase						:List[Any]						  =     type_sequence_label_size
										lowerCamelCase						:List[Any]						  =     initializer_range
										lowerCamelCase						:Any						  =     mask_ratio
										lowerCamelCase						:Optional[Any]						  =     scope

										# in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
										lowerCamelCase						:Optional[Any]						  =     (image_size // patch_size) ** 2
										lowerCamelCase						:List[Any]						  =     (num_frames // tubelet_size) * self.num_patches_per_frame

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







							def 				snake_case  (				self     :    Any		):
										lowerCamelCase						:str						  =     floats_tensor(
										    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]		)

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

										lowerCamelCase						:Any						  =     self.get_config()

										return config, pixel_values, labels







							def 				snake_case  (				self     :    List[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=__snake_case	,					initializer_range=self.initializer_range	,					)







							def 				snake_case  (				self     :    List[Any]	,					__snake_case     :    int	,					__snake_case     :    List[Any]	,					__snake_case     :    List[Any]		):
										lowerCamelCase						:str						  =     VideoMAEModel(config=__snake_case		)
										model.to(__snake_case		)
										model.eval()
										lowerCamelCase						:Optional[int]						  =     model(__snake_case		)
										self.parent.assertEqual(result.last_hidden_state.shape	,					(self.batch_size, self.seq_length, self.hidden_size)		)







							def 				snake_case  (				self     :    Tuple	,					__snake_case     :    List[Any]	,					__snake_case     :    Any	,					__snake_case     :    List[Any]		):
										lowerCamelCase						:Any						  =     VideoMAEForPreTraining(__snake_case		)
										model.to(__snake_case		)
										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
										lowerCamelCase						:int						  =     torch.ones((self.num_masks,)		)
										lowerCamelCase						:Dict						  =     torch.cat([mask, torch.zeros(self.seq_length - mask.size(0		)		)]		)
										lowerCamelCase						:Optional[Any]						  =     mask.expand(self.batch_size	,					-1		).bool()

										lowerCamelCase						:int						  =     model(__snake_case	,					__snake_case		)
										# model only returns predictions for masked patches
										lowerCamelCase						:Optional[Any]						  =     mask.sum().item()
										lowerCamelCase						:Any						  =     3 * self.tubelet_size * self.patch_size**2
										self.parent.assertEqual(result.logits.shape	,					(self.batch_size, num_masked_patches, decoder_num_labels)		)







							def 				snake_case  (				self     :    Optional[Any]		):
										lowerCamelCase						:Tuple						  =     self.prepare_config_and_inputs()
										lowerCamelCase  ,    lowerCamelCase  ,    lowerCamelCase						:Optional[Any]						  =     config_and_inputs
										lowerCamelCase						:Optional[Any]						  =     {'''pixel_values''': pixel_values}
										return config, inputs_dict


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

							_UpperCAmelCase						    =     False
							_UpperCAmelCase						    =     False
							_UpperCAmelCase						    =     False
							_UpperCAmelCase						    =     False







							def 				snake_case  (				self     :    str		):
										lowerCamelCase						:Union[str, Any]						  =     VideoMAEModelTester(self		)
										lowerCamelCase						:Tuple						  =     ConfigTester(self	,					config_class=__snake_case	,					has_text_modality=__snake_case	,					hidden_size=37		)







							def 				snake_case  (				self     :    Tuple	,					__snake_case     :    Optional[int]	,					__snake_case     :    Union[str, Any]	,					__snake_case     :    int=False		):
										lowerCamelCase						:Optional[Any]						  =     copy.deepcopy(__snake_case		)

										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
													lowerCamelCase						:Optional[int]						  =     torch.ones((self.model_tester.num_masks,)		)
													lowerCamelCase						:Tuple						  =     torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0		)		)]		)
													lowerCamelCase						:Optional[int]						  =     mask.expand(self.model_tester.batch_size	,					-1		).bool()
													lowerCamelCase						:Tuple						  =     bool_masked_pos.to(__snake_case		)

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

										return inputs_dict







							def 				snake_case  (				self     :    Tuple		):
										self.config_tester.run_common_tests()







							@unittest.skip(reason='''VideoMAE does not use inputs_embeds'''		)
							def 				snake_case  (				self     :    Dict		):
										pass







							def 				snake_case  (				self     :    Optional[Any]		):
										lowerCamelCase  ,    lowerCamelCase						:Optional[int]						  =     self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													lowerCamelCase						:List[str]						  =     model_class(__snake_case		)
													self.assertIsInstance(model.get_input_embeddings()	,					(nn.Module)		)
													lowerCamelCase						:int						  =     model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(__snake_case	,					nn.Linear		)		)







							def 				snake_case  (				self     :    List[str]		):
										lowerCamelCase  ,    lowerCamelCase						:Tuple						  =     self.model_tester.prepare_config_and_inputs_for_common()

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

													lowerCamelCase						:Optional[Any]						  =     ['''pixel_values''']
													self.assertListEqual(arg_names[:1]	,					__snake_case		)







							def 				snake_case  (				self     :    Tuple		):
										lowerCamelCase						:List[str]						  =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__snake_case		)







							def 				snake_case  (				self     :    List[Any]		):
										lowerCamelCase						:List[str]						  =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_pretraining(*__snake_case		)







							@slow
							def 				snake_case  (				self     :    Dict		):
										for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowerCamelCase						:Dict						  =     VideoMAEModel.from_pretrained(__snake_case		)
													self.assertIsNotNone(__snake_case		)







							def 				snake_case  (				self     :    str		):
										if not self.has_attentions:
													pass

										else:
													lowerCamelCase  ,    lowerCamelCase						:int						  =     self.model_tester.prepare_config_and_inputs_for_common()
													lowerCamelCase						:Any						  =     True

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

																lowerCamelCase						:List[str]						  =     True
																lowerCamelCase						:Any						  =     False
																lowerCamelCase						:Optional[Any]						  =     True
																lowerCamelCase						:Any						  =     model_class(__snake_case		)
																model.to(__snake_case		)
																model.eval()
																with torch.no_grad():
																			lowerCamelCase						:List[Any]						  =     model(**self._prepare_for_class(__snake_case	,					__snake_case		)		)
																lowerCamelCase						:Dict						  =     outputs.attentions
																self.assertEqual(len(__snake_case		)	,					self.model_tester.num_hidden_layers		)

																# check that output_attentions also work using config
																del inputs_dict["output_attentions"]
																lowerCamelCase						:Tuple						  =     True
																lowerCamelCase						:Dict						  =     model_class(__snake_case		)
																model.to(__snake_case		)
																model.eval()
																with torch.no_grad():
																			lowerCamelCase						:Optional[int]						  =     model(**self._prepare_for_class(__snake_case	,					__snake_case		)		)
																lowerCamelCase						:List[Any]						  =     outputs.attentions
																self.assertEqual(len(__snake_case		)	,					self.model_tester.num_hidden_layers		)

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

																# Check attention is always last and order is fine
																lowerCamelCase						:List[str]						  =     True
																lowerCamelCase						:Optional[Any]						  =     True
																lowerCamelCase						:Optional[int]						  =     model_class(__snake_case		)
																model.to(__snake_case		)
																model.eval()
																with torch.no_grad():
																			lowerCamelCase						:Any						  =     model(**self._prepare_for_class(__snake_case	,					__snake_case		)		)

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

																lowerCamelCase						:Any						  =     outputs.attentions

																self.assertEqual(len(__snake_case		)	,					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 				snake_case  (				self     :    List[Any]		):
										def check_hidden_states_output(__snake_case     :    Union[str, Any]	,					__snake_case     :    Union[str, Any]	,					__snake_case     :    Optional[int]		):
													lowerCamelCase						:List[str]						  =     model_class(__snake_case		)
													model.to(__snake_case		)
													model.eval()

													with torch.no_grad():
																lowerCamelCase						:List[Any]						  =     model(**self._prepare_for_class(__snake_case	,					__snake_case		)		)

													lowerCamelCase						:List[Any]						  =     outputs.hidden_states
													lowerCamelCase						:Optional[Any]						  =     self.model_tester.num_hidden_layers + 1
													self.assertEqual(len(__snake_case		)	,					__snake_case		)

													lowerCamelCase						:Union[str, Any]						  =     self.model_tester.seq_length - self.model_tester.num_masks
													lowerCamelCase						:Optional[Any]						  =     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]	,					)

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

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

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													lowerCamelCase						:Dict						  =     True

													check_hidden_states_output(__snake_case	,					__snake_case	,					__snake_case		)







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





def  _lowerCamelCase					(   ):
			lowerCamelCase						:List[str]						  =     hf_hub_download(
			    repo_id='''hf-internal-testing/spaghetti-video'''						,			filename='''eating_spaghetti.npy'''						,			repo_type='''dataset''')
			lowerCamelCase						:List[str]						  =     np.load(a_)
			return list(a_)


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







							@cached_property
							def 				snake_case  (				self     :    str		):
										# logits were tested with a different mean and std, so we use the same here
										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 				snake_case  (				self     :    Optional[Any]		):
										lowerCamelCase						:List[str]						  =     VideoMAEForVideoClassification.from_pretrained('''MCG-NJU/videomae-base-finetuned-kinetics'''		).to(
										    __snake_case		)

										lowerCamelCase						:Optional[Any]						  =     self.default_image_processor
										lowerCamelCase						:int						  =     prepare_video()
										lowerCamelCase						:Dict						  =     image_processor(__snake_case	,					return_tensors='''pt'''		).to(__snake_case		)

										# forward pass
										with torch.no_grad():
													lowerCamelCase						:List[str]						  =     model(**__snake_case		)

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

										lowerCamelCase						:List[str]						  =     torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]		).to(__snake_case		)

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







							@slow
							def 				snake_case  (				self     :    Dict		):
										lowerCamelCase						:str						  =     VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short'''		).to(__snake_case		)

										lowerCamelCase						:str						  =     self.default_image_processor
										lowerCamelCase						:Any						  =     prepare_video()
										lowerCamelCase						:List[Any]						  =     image_processor(__snake_case	,					return_tensors='''pt'''		).to(__snake_case		)

										# add boolean mask, indicating which patches to mask
										lowerCamelCase						:int						  =     hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''	,					filename='''bool_masked_pos.pt'''		)
										lowerCamelCase						:Optional[Any]						  =     torch.load(__snake_case		)

										# forward pass
										with torch.no_grad():
													lowerCamelCase						:Union[str, Any]						  =     model(**__snake_case		)

										# verify the logits
										lowerCamelCase						:List[Any]						  =     torch.Size([1, 1408, 1536]		)
										lowerCamelCase						:List[str]						  =     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=__snake_case		)
										self.assertEqual(outputs.logits.shape	,					__snake_case		)
										self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]	,					__snake_case	,					atol=1e-4		)		)

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

										# verify the loss (`config.norm_pix_loss` = `False`)
										lowerCamelCase						:List[Any]						  =     VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short'''	,					norm_pix_loss=__snake_case		).to(
										    __snake_case		)

										with torch.no_grad():
													lowerCamelCase						:Optional[int]						  =     model(**__snake_case		)

										lowerCamelCase						:Any						  =     torch.tensor(torch.tensor([0.6_4_6_9]		)	,					device=__snake_case		)
										self.assertTrue(torch.allclose(outputs.loss	,					__snake_case	,					atol=1e-4		)		)




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A__           =							{
    """configuration_nllb_moe""": [
        """NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """NllbMoeConfig""",
    ]
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			A__           =							[
			    """NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST""",
			    """NllbMoeForConditionalGeneration""",
			    """NllbMoeModel""",
			    """NllbMoePreTrainedModel""",
			    """NllbMoeTop2Router""",
			    """NllbMoeSparseMLP""",
			]


if TYPE_CHECKING:
			from .configuration_nllb_moe import (
			    NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP,
			    NllbMoeConfig,
			)

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_nllb_moe import (
						    NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST,
						    NllbMoeForConditionalGeneration,
						    NllbMoeModel,
						    NllbMoePreTrainedModel,
						    NllbMoeSparseMLP,
						    NllbMoeTopaRouter,
						)


else:
			import sys

			A__           =							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



