



import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VivitConfig
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


if is_torch_available():
      import torch
      from torch import nn

      from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VivitForVideoClassification, VivitModel
      from transformers.models.vivit.modeling_vivit import VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from transformers import VivitImageProcessor






class 				SCREAMING_SNAKE_CASE     :





       def __init__(				self						:						List[Any] ,			__lowercase						:						Any ,			__lowercase						:						Optional[int]=2 ,			__lowercase						:						List[str]=True ,			__lowercase						:						Optional[Any]=True ,			__lowercase						:						Optional[Any]=10 ,			__lowercase						:						Optional[int]=10 ,			__lowercase						:						Any=8 ,			__lowercase						:						Dict=[2, 4, 4] ,			__lowercase						:						Optional[Any]=3 ,			__lowercase						:						List[Any]=768 ,			__lowercase						:						Any=5 ,			__lowercase						:						int=4 ,			__lowercase						:						str=37 ,			__lowercase						:						str="gelu_fast" ,			__lowercase						:						Any=0.0 ,			__lowercase						:						List[str]=0.0 ,			__lowercase						:						Union[str, Any]=0.02 ,			__lowercase						:						Optional[int]=1E-06 ,			__lowercase						:						Dict=True ,			__lowercase						:						List[Any]=None ,			):

           '''simple docstring'''
           __a					       =						parent
           __a					       =						batch_size
           __a					       =						is_training
           __a					       =						use_labels
           __a					       =						num_labels
           __a					       =						image_size
           __a					       =						num_frames
           __a					       =						tubelet_size
           __a					       =						num_channels
           __a					       =						hidden_size
           __a					       =						num_hidden_layers
           __a					       =						num_attention_heads
           __a					       =						intermediate_size
           __a					       =						hidden_act
           __a					       =						hidden_dropout_prob
           __a					       =						attention_probs_dropout_prob
           __a					       =						initializer_range
           __a					       =						layer_norm_eps
           __a					       =						qkv_bias
           __a					       =						scope

           __a					       =						(
               (self.image_size // self.tubelet_size[2])
               * (self.image_size // self.tubelet_size[1])
               * (self.num_frames // self.tubelet_size[0])
           ) + 1  # CLS token





       def   UpperCamelCase_   (				self						:						List[str]		):

           '''simple docstring'''
           __a					       =						floats_tensor(
               [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]		)

           __a					       =						None
           if self.use_labels:
               __a					       =						ids_tensor([self.batch_size] ,			self.num_labels		)

           __a					       =						self.get_config()

           return config, pixel_values, labels





       def   UpperCamelCase_   (				self						:						str		):

           '''simple docstring'''
           __a					       =						VivitConfig(
               num_frames=self.num_frames ,			image_size=self.image_size ,			tubelet_size=self.tubelet_size ,			num_channels=self.num_channels ,			hidden_size=self.hidden_size ,			num_hidden_layers=self.num_hidden_layers ,			num_attention_heads=self.num_attention_heads ,			intermediate_size=self.intermediate_size ,			hidden_act=self.hidden_act ,			hidden_dropout_prob=self.hidden_dropout_prob ,			attention_probs_dropout_prob=self.attention_probs_dropout_prob ,			initializer_range=self.initializer_range ,			layer_norm_eps=self.layer_norm_eps ,			qkv_bias=self.qkv_bias ,			)
           __a					       =						self.num_labels
           return config





       def   UpperCamelCase_   (				self						:						Tuple ,			__lowercase						:						Union[str, Any] ,			__lowercase						:						Any ,			__lowercase						:						Any		):

           '''simple docstring'''
           __a					       =						VivitModel(config=__lowercase		)
           model.to(__lowercase		)
           model.eval()
           __a					       =						model(__lowercase		)
           self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, self.seq_length, self.hidden_size)		)





       def   UpperCamelCase_   (				self						:						str ,			__lowercase						:						int ,			__lowercase						:						List[str] ,			__lowercase						:						Optional[Any]		):

           '''simple docstring'''
           __a					       =						VivitForVideoClassification(__lowercase		)
           model.to(__lowercase		)
           model.eval()

           __a					       =						model(__lowercase		)

           # verify the logits shape
           __a					       =						torch.Size((self.batch_size, self.num_labels)		)
           self.parent.assertEqual(result.logits.shape ,			__lowercase		)







       def   UpperCamelCase_   (				self						:						Optional[Any]		):

           '''simple docstring'''
           __a					       =						self.prepare_config_and_inputs()
           __a       ,   __a       ,   __a					       =						config_and_inputs
           __a					       =						{"""pixel_values""": pixel_values}
           return config, inputs_dict






@require_torch
class 				SCREAMING_SNAKE_CASE     (							lowerCamelCase__     , unittest.TestCase   ):
       __lowerCamelCase		:		List[Any]      =(VivitModel, VivitForVideoClassification) if is_torch_available() else ()

       __lowerCamelCase		:		Dict      =False
       __lowerCamelCase		:		Dict      =False
       __lowerCamelCase		:		Optional[Any]      =False
       __lowerCamelCase		:		Dict      =False





       def   UpperCamelCase_   (				self						:						Any		):

           '''simple docstring'''
           __a					       =						VivitModelTester(self		)
           __a					       =						ConfigTester(self ,			config_class=__lowercase ,			has_text_modality=__lowercase ,			hidden_size=37		)





       def   UpperCamelCase_   (				self						:						Optional[Any] ,			__lowercase						:						List[Any] ,			__lowercase						:						Union[str, Any] ,			__lowercase						:						List[Any]=False		):

           '''simple docstring'''
           __a					       =						copy.deepcopy(__lowercase		)

           if return_labels:
               if model_class in get_values(__lowercase		):
                   __a					       =						torch.zeros(
                       self.model_tester.batch_size ,			dtype=torch.long ,			device=__lowercase		)

           return inputs_dict





       def   UpperCamelCase_   (				self						:						Dict		):

           '''simple docstring'''
           self.config_tester.run_common_tests()





       @unittest.skip(reason="""Vivit does not use inputs_embeds"""		)
       def   UpperCamelCase_   (				self						:						List[str]		):

           '''simple docstring'''
           pass





       def   UpperCamelCase_   (				self						:						List[str]		):

           '''simple docstring'''
           __a       ,   __a					       =						self.model_tester.prepare_config_and_inputs_for_common()

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





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

           '''simple docstring'''
           __a       ,   __a					       =						self.model_tester.prepare_config_and_inputs_for_common()

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

               __a					       =						["""pixel_values""", """head_mask"""]
               self.assertListEqual(arg_names[:2] ,			__lowercase		)





       def   UpperCamelCase_   (				self						:						List[str]		):

           '''simple docstring'''
           __a					       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*__lowercase		)





       def   UpperCamelCase_   (				self						:						Optional[int]		):

           '''simple docstring'''
           __a					       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_video_classification(*__lowercase		)





       @slow
       def   UpperCamelCase_   (				self						:						str		):

           '''simple docstring'''
           for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __a					       =						VivitModel.from_pretrained(__lowercase		)
               self.assertIsNotNone(__lowercase		)





       def   UpperCamelCase_   (				self						:						Tuple		):

           '''simple docstring'''
           __a       ,   __a					       =						self.model_tester.prepare_config_and_inputs_for_common()
           __a					       =						True

           for model_class in self.all_model_classes:
               __a					       =						self.model_tester.seq_length

               __a					       =						True
               __a					       =						False
               __a					       =						True
               __a					       =						model_class(__lowercase		)
               model.to(__lowercase		)
               model.eval()
               with torch.no_grad():
                   __a					       =						model(**self._prepare_for_class(__lowercase ,			__lowercase		)		)
               __a					       =						outputs.attentions
               self.assertEqual(len(__lowercase		) ,			self.model_tester.num_hidden_layers		)

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

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

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

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

               __a					       =						outputs.attentions

               self.assertEqual(len(__lowercase		) ,			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						:						Optional[int]		):

           '''simple docstring'''
           def check_hidden_states_output(__lowercase						:						Optional[int] ,			__lowercase						:						str ,			__lowercase						:						Union[str, Any]		):
               __a					       =						model_class(__lowercase		)
               model.to(__lowercase		)
               model.eval()

               with torch.no_grad():
                   __a					       =						model(**self._prepare_for_class(__lowercase ,			__lowercase		)		)

               __a					       =						outputs.hidden_states
               __a					       =						self.model_tester.num_hidden_layers + 1
               self.assertEqual(len(__lowercase		) ,			__lowercase		)

               __a					       =						self.model_tester.seq_length

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

           __a       ,   __a					       =						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               __a					       =						True
               check_hidden_states_output(__lowercase ,			__lowercase ,			__lowercase		)

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

               check_hidden_states_output(__lowercase ,			__lowercase ,			__lowercase		)







       @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""		)
       def   UpperCamelCase_   (				self						:						List[Any]		):

           '''simple docstring'''
           pass



def 					lowerCAmelCase__  (							):
    """simple docstring"""







    __a					       =						hf_hub_download(
        repo_id="""hf-internal-testing/spaghetti-video"""			,       filename="""eating_spaghetti_32_frames.npy"""			,       repo_type="""dataset"""      )
    __a					       =						np.load(_SCREAMING_SNAKE_CASE      )
    return list(_SCREAMING_SNAKE_CASE      )






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





       @cached_property
       def   UpperCamelCase_   (				self						:						Tuple		):

           '''simple docstring'''
           return VivitImageProcessor() if is_vision_available() else None







       @slow
       def   UpperCamelCase_   (				self						:						Dict		):

           '''simple docstring'''
           __a					       =						VivitForVideoClassification.from_pretrained("""google/vivit-b-16x2-kinetics400"""		).to(__lowercase		)

           __a					       =						self.default_image_processor
           __a					       =						prepare_video()
           __a					       =						image_processor(__lowercase ,			return_tensors="""pt"""		).to(__lowercase		)

           # forward pass
           with torch.no_grad():
               __a					       =						model(**__lowercase		)

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

           # taken from original model
           __a					       =						torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]		).to(__lowercase		)

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





import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json

def 					lowerCAmelCase__  (							_SCREAMING_SNAKE_CASE    :     Dict      ):
				"""simple docstring"""







				__a					       =						filter(lambda _SCREAMING_SNAKE_CASE      : p.requires_grad			,       model.parameters()      )
				__a					       =						sum([np.prod(p.size()      ) for p in model_parameters]      )
				return params


lowerCamelCase__     =	logging.getLogger(__name__)






class 				SCREAMING_SNAKE_CASE     (							pl.Callback   ):





							def   UpperCamelCase_   (				self						:						Dict ,			__lowercase						:						Dict ,			__lowercase						:						str		):

											'''simple docstring'''
											__a					       =						{F"lr_group_{i}": param["""lr"""] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups		)}
											pl_module.logger.log_metrics(__lowercase		)





							@rank_zero_only
							def   UpperCamelCase_   (				self						:						Optional[Any] ,			__lowercase						:						pl.Trainer ,			__lowercase						:						pl.LightningModule ,			__lowercase						:						str ,			__lowercase						:						Tuple=True		):

											'''simple docstring'''
											logger.info(F"***** {type_path} results at step {trainer.global_step:05d} *****"		)
											__a					       =						trainer.callback_metrics
											trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["""log""", """progress_bar""", """preds"""]}		)
											# Log results
											__a					       =						Path(pl_module.hparams.output_dir		)
											if type_path == "test":
															__a					       =						od / """test_results.txt"""
															__a					       =						od / """test_generations.txt"""
											else:
															# this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
															# If people want this it will be easy enough to add back.
															__a					       =						od / F"{type_path}_results/{trainer.global_step:05d}.txt"
															__a					       =						od / F"{type_path}_generations/{trainer.global_step:05d}.txt"
															results_file.parent.mkdir(exist_ok=__lowercase		)
															generations_file.parent.mkdir(exist_ok=__lowercase		)
											with open(__lowercase ,			"""a+"""		) as writer:
															for key in sorted(__lowercase		):
																			if key in ["log", "progress_bar", "preds"]:
																							continue
																			__a					       =						metrics[key]
																			if isinstance(__lowercase ,			torch.Tensor		):
																							__a					       =						val.item()
																			__a					       =						F"{key}: {val:.6f}\n"
																			writer.write(__lowercase		)

											if not save_generations:
															return

											if "preds" in metrics:
															__a					       =						"""\n""".join(metrics["""preds"""]		)
															generations_file.open("""w+"""		).write(__lowercase		)





							@rank_zero_only
							def   UpperCamelCase_   (				self						:						int ,			__lowercase						:						Union[str, Any] ,			__lowercase						:						List[Any]		):

											'''simple docstring'''
											try:
															__a					       =						pl_module.model.model.num_parameters()
											except AttributeError:
															__a					       =						pl_module.model.num_parameters()

											__a					       =						count_trainable_parameters(__lowercase		)
											# mp stands for million parameters
											trainer.logger.log_metrics({"""n_params""": npars, """mp""": npars / 1E6, """grad_mp""": n_trainable_pars / 1E6}		)





							@rank_zero_only
							def   UpperCamelCase_   (				self						:						Union[str, Any] ,			__lowercase						:						pl.Trainer ,			__lowercase						:						pl.LightningModule		):

											'''simple docstring'''
											save_json(pl_module.metrics ,			pl_module.metrics_save_path		)
											return self._write_logs(__lowercase ,			__lowercase ,			"""test"""		)







							@rank_zero_only
							def   UpperCamelCase_   (				self						:						Optional[int] ,			__lowercase						:						pl.Trainer ,			__lowercase						:						Dict		):

											'''simple docstring'''
											save_json(pl_module.metrics ,			pl_module.metrics_save_path		)
											# Uncommenting this will save val generations
											# return self._write_logs(trainer, pl_module, "valid")



def 					lowerCAmelCase__  (							_SCREAMING_SNAKE_CASE    :     Optional[int]			,       _SCREAMING_SNAKE_CASE    :     Optional[Any]			,       _SCREAMING_SNAKE_CASE    :     Tuple=1			,       _SCREAMING_SNAKE_CASE    :     List[str]=False      ):
				"""simple docstring"""







				if metric == "rouge2":
								__a					       =						"""{val_avg_rouge2:.4f}-{step_count}"""
				elif metric == "bleu":
								__a					       =						"""{val_avg_bleu:.4f}-{step_count}"""
				elif metric == "loss":
								__a					       =						"""{val_avg_loss:.4f}-{step_count}"""
				else:
								raise NotImplementedError(
								    f"seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"
								    """ this function."""      )

				__a					       =						ModelCheckpoint(
				    dirpath=_SCREAMING_SNAKE_CASE			,       filename=_SCREAMING_SNAKE_CASE			,       monitor=f"val_{metric}"			,       mode="""min""" if """loss""" in metric else """max"""			,       save_top_k=_SCREAMING_SNAKE_CASE			,       )
				return checkpoint_callback

def 					lowerCAmelCase__  (							_SCREAMING_SNAKE_CASE    :     int			,       _SCREAMING_SNAKE_CASE    :     Optional[Any]      ):
				"""simple docstring"""







				return EarlyStopping(
				    monitor=f"val_{metric}"			,       mode="""min""" if """loss""" in metric else """max"""			,       patience=_SCREAMING_SNAKE_CASE			,       verbose=_SCREAMING_SNAKE_CASE			,       )
