
import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

from seqaseq_trainer import SeqaSeqTrainer
from seqaseq_training_args import SeqaSeqTrainingArguments

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSeqaSeqLM,
    AutoTokenizer,
    HfArgumentParser,
    MBartTokenizer,
    MBartTokenizerFast,
    set_seed,
)
from transformers.trainer_utils import EvaluationStrategy, is_main_process
from transformers.training_args import ParallelMode
from utils import (
    SeqaSeqDataCollator,
    SeqaSeqDataset,
    assert_all_frozen,
    build_compute_metrics_fn,
    check_output_dir,
    freeze_embeds,
    freeze_params,
    lmap,
    save_json,
    use_task_specific_params,
    write_txt_file,
)


__a         =   logging.getLogger(__name__)
@dataclass
class 					lowercase__:



	"""simple docstring"""




	a     :str       	=      field(
	    metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}				)
	a     :Optional[str]       	=      field(
	    default=UpperCAmelCase					,		metadata={'help': 'Pretrained config name or path if not the same as model_name'}				)
	a     :Optional[str]       	=      field(
	    default=UpperCAmelCase					,		metadata={'help': 'Pretrained tokenizer name or path if not the same as model_name'}				)
	a     :Optional[str]       	=      field(
	    default=UpperCAmelCase					,		metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}					,		)
	a     :bool       	=      field(default=UpperCAmelCase					,		metadata={'help': 'Whether tp freeze the encoder.'}				)
	a     :bool       	=      field(default=UpperCAmelCase					,		metadata={'help': 'Whether  to freeze the embeddings.'}				)

@dataclass
class 					lowercase__:



	"""simple docstring"""




	a     :str       	=      field(
	    metadata={'help': 'The input data dir. Should contain the .tsv files (or other data files) for the task.'}				)
	a     :Optional[str]       	=      field(
	    default='summarization'					,		metadata={'help': 'Task name, summarization (or summarization_{dataset} for pegasus) or translation'}					,		)
	a     :Optional[int]       	=      field(
	    default=1_024					,		metadata={
	        'help': (
	            'The maximum total input sequence length after tokenization. Sequences longer '
	            'than this will be truncated, sequences shorter will be padded.'
	        )
	    }					,		)
	a     :Optional[int]       	=      field(
	    default=128					,		metadata={
	        'help': (
	            'The maximum total sequence length for target text after tokenization. Sequences longer '
	            'than this will be truncated, sequences shorter will be padded.'
	        )
	    }					,		)
	a     :Optional[int]       	=      field(
	    default=142					,		metadata={
	        'help': (
	            'The maximum total sequence length for validation target text after tokenization. Sequences longer '
	            'than this will be truncated, sequences shorter will be padded. '
	            'This argument is also used to override the ``max_length`` param of ``model.generate``, which is used '
	            'during ``evaluate`` and ``predict``.'
	        )
	    }					,		)
	a     :Optional[int]       	=      field(
	    default=142					,		metadata={
	        'help': (
	            'The maximum total sequence length for test target text after tokenization. Sequences longer '
	            'than this will be truncated, sequences shorter will be padded.'
	        )
	    }					,		)
	a     :Optional[int]       	=      field(default=-1					,		metadata={'help': '# training examples. -1 means use all.'}				)
	a     :Optional[int]       	=      field(default=-1					,		metadata={'help': '# validation examples. -1 means use all.'}				)
	a     :Optional[int]       	=      field(default=-1					,		metadata={'help': '# test examples. -1 means use all.'}				)
	a     :Optional[str]       	=      field(default=UpperCAmelCase					,		metadata={'help': 'Source language id for translation.'}				)
	a     :Optional[str]       	=      field(default=UpperCAmelCase					,		metadata={'help': 'Target language id for translation.'}				)
	a     :Optional[int]       	=      field(default=UpperCAmelCase					,		metadata={'help': '# num_beams to use for evaluation.'}				)
	a     :bool       	=      field(
	    default=UpperCAmelCase					,		metadata={'help': 'If only pad tokens should be ignored. This assumes that `config.pad_token_id` is defined.'}					,		)




def 							a		(				snake_case__:   Dict		,  snake_case__:   Optional[int]		,  snake_case__:   List[str]   ):


				'''simple docstring'''
				logger.info(F'''***** {split} metrics *****'''   )
				for key in sorted(metrics.keys()   ):
								logger.info(F'''  {key} = {metrics[key]}'''   )
				save_json(snake_case__		,  os.path.join(snake_case__		,  F'''{split}_results.json'''   )   )




def 							a		(				):


				'''simple docstring'''
				# See all possible arguments in src/transformers/training_args.py
				# or by passing the --help flag to this script.
				# We now keep distinct sets of args, for a cleaner separation of concerns.

				lowercase_	    =   HfArgumentParser((ModelArguments, DataTrainingArguments, SeqaSeqTrainingArguments)   )

				if len(sys.argv   ) == 2 and sys.argv[1].endswith('''.json'''   ):
								# If we pass only one argument to the script and it's the path to a json file,
								# let's parse it to get our arguments.
								lowercase_    ,      lowercase_    ,      lowercase_	    =   parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]   )   )
				else:
								lowercase_    ,      lowercase_    ,      lowercase_	    =   parser.parse_args_into_dataclasses()

				check_output_dir(snake_case__   )

				# Setup logging
				logging.basicConfig(
				    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''		,  datefmt='''%m/%d/%Y %H:%M:%S'''		,  level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN		,  )
				logger.warning(
				    '''Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s'''		,  training_args.local_rank		,  training_args.device		,  training_args.n_gpu		,  bool(training_args.parallel_mode == ParallelMode.DISTRIBUTED   )		,  training_args.fpaa		,  )
				transformers.utils.logging.enable_default_handler()
				transformers.utils.logging.enable_explicit_format()
				# Set the verbosity to info of the Transformers logger (on main process only):
				if is_main_process(training_args.local_rank   ):
								transformers.utils.logging.set_verbosity_info()
				logger.info('''Training/evaluation parameters %s'''		,  snake_case__   )

				# Set seed
				set_seed(training_args.seed   )

				# Load pretrained model and tokenizer
				#
				# Distributed training:
				# The .from_pretrained methods guarantee that only one local process can concurrently
				# download model & vocab.

				lowercase_	    =   AutoConfig.from_pretrained(
				    model_args.config_name if model_args.config_name else model_args.model_name_or_path		,  cache_dir=model_args.cache_dir		,  )

				lowercase_	    =   ('''encoder_layerdrop''', '''decoder_layerdrop''', '''dropout''', '''attention_dropout''')
				for p in extra_model_params:
								if getattr(snake_case__		,  snake_case__		,  snake_case__   ):
												assert hasattr(snake_case__		,  snake_case__   ), F'''({config.__class__.__name__}) doesn\'t have a `{p}` attribute'''
												setattr(snake_case__		,  snake_case__		,  getattr(snake_case__		,  snake_case__   )   )

				lowercase_	    =   AutoTokenizer.from_pretrained(
				    model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path		,  cache_dir=model_args.cache_dir		,  )
				lowercase_	    =   AutoModelForSeqaSeqLM.from_pretrained(
				    model_args.model_name_or_path		,  from_tf='''.ckpt''' in model_args.model_name_or_path		,  config=snake_case__		,  cache_dir=model_args.cache_dir		,  )

				# use task specific params
				use_task_specific_params(snake_case__		,  data_args.task   )

				# set num_beams for evaluation
				if data_args.eval_beams is None:
								lowercase_	    =   model.config.num_beams

				# set decoder_start_token_id for MBart
				if model.config.decoder_start_token_id is None and isinstance(snake_case__		,  (MBartTokenizer, MBartTokenizerFast)   ):
								assert (
								    data_args.tgt_lang is not None and data_args.src_lang is not None
								), "mBart requires --tgt_lang and --src_lang"
								if isinstance(snake_case__		,  snake_case__   ):
												lowercase_	    =   tokenizer.lang_code_to_id[data_args.tgt_lang]
								else:
												lowercase_	    =   tokenizer.convert_tokens_to_ids(data_args.tgt_lang   )

				if model_args.freeze_embeds:
								freeze_embeds(snake_case__   )
				if model_args.freeze_encoder:
								freeze_params(model.get_encoder()   )
								assert_all_frozen(model.get_encoder()   )

				lowercase_	    =   SeqaSeqDataset

				# Get datasets
				lowercase_	    =   (
				    dataset_class(
				        snake_case__		,  type_path='''train'''		,  data_dir=data_args.data_dir		,  n_obs=data_args.n_train		,  max_target_length=data_args.max_target_length		,  max_source_length=data_args.max_source_length		,  prefix=model.config.prefix or ''''''		,  )
				    if training_args.do_train
				    else None
				)
				lowercase_	    =   (
				    dataset_class(
				        snake_case__		,  type_path='''val'''		,  data_dir=data_args.data_dir		,  n_obs=data_args.n_val		,  max_target_length=data_args.val_max_target_length		,  max_source_length=data_args.max_source_length		,  prefix=model.config.prefix or ''''''		,  )
				    if training_args.do_eval or training_args.evaluation_strategy != EvaluationStrategy.NO
				    else None
				)
				lowercase_	    =   (
				    dataset_class(
				        snake_case__		,  type_path='''test'''		,  data_dir=data_args.data_dir		,  n_obs=data_args.n_test		,  max_target_length=data_args.test_max_target_length		,  max_source_length=data_args.max_source_length		,  prefix=model.config.prefix or ''''''		,  )
				    if training_args.do_predict
				    else None
				)

				# Initialize our Trainer
				lowercase_	    =   (
				    build_compute_metrics_fn(data_args.task		,  snake_case__   ) if training_args.predict_with_generate else None
				)
				lowercase_	    =   SeqaSeqTrainer(
				    model=snake_case__		,  args=snake_case__		,  data_args=snake_case__		,  train_dataset=snake_case__		,  eval_dataset=snake_case__		,  data_collator=SeqaSeqDataCollator(
				        snake_case__		,  snake_case__		,  model.config.decoder_start_token_id		,  training_args.tpu_num_cores   )		,  compute_metrics=snake_case__		,  tokenizer=snake_case__		,  )

				lowercase_	    =   {}
				# Training
				if training_args.do_train:
								logger.info('''*** Train ***'''   )

								lowercase_	    =   trainer.train(
								    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path   ) else None   )
								lowercase_	    =   train_result.metrics
								lowercase_	    =   data_args.n_train

								trainer.save_model()  # this also saves the tokenizer

								if trainer.is_world_process_zero():
												handle_metrics('''train'''		,  snake_case__		,  training_args.output_dir   )
												all_metrics.update(snake_case__   )

												# Need to save the state, since Trainer.save_model saves only the tokenizer with the model
												trainer.state.save_to_json(os.path.join(training_args.output_dir		,  '''trainer_state.json'''   )   )

												# For convenience, we also re-save the tokenizer to the same directory,
												# so that you can share your model easily on huggingface.co/models =)
												tokenizer.save_pretrained(training_args.output_dir   )

    # Evaluation
				if training_args.do_eval:
								logger.info('''*** Evaluate ***'''   )

								lowercase_	    =   trainer.evaluate(metric_key_prefix='''val'''   )
								lowercase_	    =   data_args.n_val
								lowercase_	    =   round(metrics['''val_loss''']		,  4   )

								if trainer.is_world_process_zero():
												handle_metrics('''val'''		,  snake_case__		,  training_args.output_dir   )
												all_metrics.update(snake_case__   )

				if training_args.do_predict:
								logger.info('''*** Predict ***'''   )

								lowercase_	    =   trainer.predict(test_dataset=snake_case__		,  metric_key_prefix='''test'''   )
								lowercase_	    =   test_output.metrics
								lowercase_	    =   data_args.n_test

								if trainer.is_world_process_zero():
												lowercase_	    =   round(metrics['''test_loss''']		,  4   )
												handle_metrics('''test'''		,  snake_case__		,  training_args.output_dir   )
												all_metrics.update(snake_case__   )

												if training_args.predict_with_generate:
																lowercase_	    =   tokenizer.batch_decode(
																    test_output.predictions		,  skip_special_tokens=snake_case__		,  clean_up_tokenization_spaces=snake_case__   )
																lowercase_	    =   lmap(str.strip		,  snake_case__   )
																write_txt_file(snake_case__		,  os.path.join(training_args.output_dir		,  '''test_generations.txt'''   )   )

				if trainer.is_world_process_zero():
								save_json(snake_case__		,  os.path.join(training_args.output_dir		,  '''all_results.json'''   )   )

				return all_metrics




def 							a		(				snake_case__:   List[str]   ):


				'''simple docstring'''
				# For xla_spawn (TPUs)
				main()


if __name__ == "__main__":
						main()









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



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

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

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



	def 		UpperCamelCase				(    self				:     Any)->				int:
						__lowerCAmelCase						 =floats_tensor(
						    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size])

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

						__lowerCAmelCase						 =self.get_config()

						return config, pixel_values, labels



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



	def 		UpperCamelCase				(    self				:     Any					,						snake_case_				:     Union[str, Any]					,						snake_case_				:     List[Any]					,						snake_case_				:     Tuple)->				Optional[int]:
						__lowerCAmelCase						 =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						 =torch.ones((self.num_masks,))
						__lowerCAmelCase						 =torch.cat([mask, torch.zeros(self.seq_length - mask.size(0))])
						__lowerCAmelCase						 =mask.expand(self.batch_size					,						-1).bool()

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



	def 		UpperCamelCase				(    self				:     List[Any])->				List[str]:
						__lowerCAmelCase						 =self.prepare_config_and_inputs()
						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase						 =config_and_inputs
						__lowerCAmelCase						 ={"""pixel_values""": pixel_values}
						return config, inputs_dict






@require_torch
class     __a			(  SCREAMING_SNAKE_CASE						,     SCREAMING_SNAKE_CASE						,     unittest.TestCase					):
	SCREAMING_SNAKE_CASE  	=							(
	    (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
	)
	SCREAMING_SNAKE_CASE  	=							(
	    {"feature-extraction": VideoMAEModel, "video-classification": VideoMAEForVideoClassification}
	    if is_torch_available()
	    else {}
	)

	SCREAMING_SNAKE_CASE  	=							False
	SCREAMING_SNAKE_CASE  	=							False
	SCREAMING_SNAKE_CASE  	=							False
	SCREAMING_SNAKE_CASE  	=							False



	def 		UpperCamelCase				(    self				:     List[Any])->				List[Any]:
						__lowerCAmelCase						 =VideoMAEModelTester(self)
						__lowerCAmelCase						 =ConfigTester(self					,						config_class=snake_case_					,						has_text_modality=snake_case_					,						hidden_size=37)



	def 		UpperCamelCase				(    self				:     Tuple					,						snake_case_				:     str					,						snake_case_				:     Tuple					,						snake_case_				:     List[str]=False)->				Tuple:
						__lowerCAmelCase						 =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						 =torch.ones((self.model_tester.num_masks,))
											__lowerCAmelCase						 =torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0))])
											__lowerCAmelCase						 =mask.expand(self.model_tester.batch_size					,						-1).bool()
											__lowerCAmelCase						 =bool_masked_pos.to(snake_case_)

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

						return inputs_dict



	def 		UpperCamelCase				(    self				:     List[str])->				Any:
						self.config_tester.run_common_tests()



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



	def 		UpperCamelCase				(    self				:     Union[str, Any])->				Tuple:
						__lowerCAmelCase ,						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs_for_common()

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



	def 		UpperCamelCase				(    self				:     Optional[int])->				List[str]:
						__lowerCAmelCase ,						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs_for_common()

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

											__lowerCAmelCase						 =["""pixel_values"""]
											self.assertListEqual(arg_names[:1]					,						snake_case_)



	def 		UpperCamelCase				(    self				:     str)->				List[Any]:
						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*snake_case_)



	def 		UpperCamelCase				(    self				:     List[Any])->				Any:
						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*snake_case_)



	@slow
	def 		UpperCamelCase				(    self				:     List[Any])->				Any:
						for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__lowerCAmelCase						 =VideoMAEModel.from_pretrained(snake_case_)
											self.assertIsNotNone(snake_case_)



	def 		UpperCamelCase				(    self				:     List[Any])->				List[str]:
						if not self.has_attentions:
											pass

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

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

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

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

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

																__lowerCAmelCase						 =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 		UpperCamelCase				(    self				:     Tuple)->				Any:
						def check_hidden_states_output(snake_case_				:     Tuple					,						snake_case_				:     Tuple					,						snake_case_				:     int):
											__lowerCAmelCase						 =model_class(snake_case_)
											model.to(snake_case_)
											model.eval()

											with torch.no_grad():
																__lowerCAmelCase						 =model(**self._prepare_for_class(snake_case_					,						snake_case_))

											__lowerCAmelCase						 =outputs.hidden_states
											__lowerCAmelCase						 =self.model_tester.num_hidden_layers + 1
											self.assertEqual(len(snake_case_)					,						snake_case_)

											__lowerCAmelCase						 =self.model_tester.seq_length - self.model_tester.num_masks
											__lowerCAmelCase						 =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						 =self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
											__lowerCAmelCase						 =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						 =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 		UpperCamelCase				(    self				:     Optional[Any])->				int:
						pass






def        __lowerCAmelCase		(	)	->    Any:
					__lowerCAmelCase						 =hf_hub_download(
					    repo_id="""hf-internal-testing/spaghetti-video"""					, filename="""eating_spaghetti.npy"""					, repo_type="""dataset"""						)
					__lowerCAmelCase						 =np.load(__lowerCamelCase						)
					return list(__lowerCamelCase						)






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



	@cached_property
	def 		UpperCamelCase				(    self				:     str)->				List[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 		UpperCamelCase				(    self				:     Optional[Any])->				Tuple:
						__lowerCAmelCase						 =VideoMAEForVideoClassification.from_pretrained("""MCG-NJU/videomae-base-finetuned-kinetics""").to(
						    snake_case_)

						__lowerCAmelCase						 =self.default_image_processor
						__lowerCAmelCase						 =prepare_video()
						__lowerCAmelCase						 =image_processor(snake_case_					,						return_tensors="""pt""").to(snake_case_)

						# forward pass
						with torch.no_grad():
											__lowerCAmelCase						 =model(**snake_case_)

						# verify the logits
						__lowerCAmelCase						 =torch.Size((1, 4_00))
						self.assertEqual(outputs.logits.shape					,						snake_case_)

						__lowerCAmelCase						 =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 		UpperCamelCase				(    self				:     str)->				str:
						__lowerCAmelCase						 =VideoMAEForPreTraining.from_pretrained("""MCG-NJU/videomae-base-short""").to(snake_case_)

						__lowerCAmelCase						 =self.default_image_processor
						__lowerCAmelCase						 =prepare_video()
						__lowerCAmelCase						 =image_processor(snake_case_					,						return_tensors="""pt""").to(snake_case_)

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

						# forward pass
						with torch.no_grad():
											__lowerCAmelCase						 =model(**snake_case_)

						# verify the logits
						__lowerCAmelCase						 =torch.Size([1, 14_08, 15_36])
						__lowerCAmelCase						 =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						 =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						 =VideoMAEForPreTraining.from_pretrained("""MCG-NJU/videomae-base-short"""					,						norm_pix_loss=snake_case_).to(
						    snake_case_)

						with torch.no_grad():
											__lowerCAmelCase						 =model(**snake_case_)

						__lowerCAmelCase						 =torch.tensor(torch.tensor([0.6_4_6_9])					,						device=snake_case_)
						self.assertTrue(torch.allclose(outputs.loss					,						snake_case_					,						atol=1e-4))




