import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
 import jax
 import jax.numpy as jnp

 from transformers.modeling_flax_pytorch_utils import (
     convert_pytorch_state_dict_to_flax,
     load_flax_weights_in_pytorch_model,
 )
 from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
 import torch




class  __lowerCAmelCase	:



 '''simple docstring'''

 def __init__(      self							:		int							,      lowerCamelCase							:		Union[str, Any]							,      lowerCamelCase							:		Dict=1_4							,      lowerCamelCase							:		Dict=7							,      lowerCamelCase							:		str=True							,      lowerCamelCase							:		List[Any]=True							,      lowerCamelCase							:		Any=False							,      lowerCamelCase							:		int=True							,      lowerCamelCase							:		List[str]=9_9							,      lowerCamelCase							:		Dict=3_2							,      lowerCamelCase							:		Any=5							,      lowerCamelCase							:		List[Any]=4							,      lowerCamelCase							:		Tuple=3_7							,      lowerCamelCase							:		Tuple="gelu"							,      lowerCamelCase							:		int=0.1							,      lowerCamelCase							:		List[str]=0.1							,      lowerCamelCase							:		int=5_1_2							,      lowerCamelCase							:		Union[str, Any]=0.0_2							,      )			->				Optional[int]:
        _UpperCamelCase  :							Any	    =	parent
        _UpperCamelCase  :							List[Any]	    =	batch_size
        _UpperCamelCase  :							List[str]	    =	seq_length
        _UpperCamelCase  :							Any	    =	is_training
        _UpperCamelCase  :							Optional[int]	    =	use_input_mask
        _UpperCamelCase  :							int	    =	use_token_type_ids
        _UpperCamelCase  :							Any	    =	use_labels
        _UpperCamelCase  :							str	    =	vocab_size
        _UpperCamelCase  :							Any	    =	hidden_size
        _UpperCamelCase  :							Any	    =	num_hidden_layers
        _UpperCamelCase  :							int	    =	num_attention_heads
        _UpperCamelCase  :							Dict	    =	intermediate_size
        _UpperCamelCase  :							Any	    =	hidden_act
        _UpperCamelCase  :							Optional[int]	    =	hidden_dropout_prob
        _UpperCamelCase  :							Any	    =	attention_probs_dropout_prob
        _UpperCamelCase  :							Tuple	    =	max_position_embeddings
        _UpperCamelCase  :							Tuple	    =	initializer_range
        _UpperCamelCase  :							Union[str, Any]	    =	None
        _UpperCamelCase  :							Dict	    =	vocab_size - 1
        _UpperCamelCase  :							Optional[int]	    =	vocab_size - 1
        _UpperCamelCase  :							Optional[Any]	    =	vocab_size - 1

 def        __snake_case       (      self							:		List[Any]	)			->				Any:
        _UpperCamelCase  :							Dict	    =	ids_tensor([self.batch_size, self.seq_length]							,      self.vocab_size	)

        _UpperCamelCase  :							Optional[int]	    =	None
        if self.use_input_mask:
               _UpperCamelCase  :							Union[str, Any]	    =	random_attention_mask([self.batch_size, self.seq_length]	)

        _UpperCamelCase  :							str	    =	GPTaConfig(
            vocab_size=self.vocab_size							,      n_embd=self.hidden_size							,      n_layer=self.num_hidden_layers							,      n_head=self.num_attention_heads							,      n_positions=self.max_position_embeddings							,      use_cache=lowerCamelCase							,      bos_token_id=self.bos_token_id							,      eos_token_id=self.eos_token_id							,      pad_token_id=self.pad_token_id							,      )

        return (config, input_ids, input_mask)

 def        __snake_case       (      self							:		int	)			->				Union[str, Any]:
        _UpperCamelCase  :							Any	    =	self.prepare_config_and_inputs()
        _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							Optional[int]	    =	config_and_inputs
        _UpperCamelCase  :							Optional[int]	    =	{'input_ids': input_ids, 'attention_mask': attention_mask}
        return config, inputs_dict

 def        __snake_case       (      self							:		Optional[int]	)			->				List[str]:
        _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							int	    =	self.prepare_config_and_inputs()

        _UpperCamelCase  :							List[str]	    =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
        _UpperCamelCase  :							str	    =	ids_tensor([self.batch_size, self.seq_length]							,      vocab_size=2	)

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )

 def        __snake_case       (      self							:		str							,      lowerCamelCase							:		Union[str, Any]							,      lowerCamelCase							:		str							,      lowerCamelCase							:		List[Any]							,      lowerCamelCase							:		Optional[int]	)			->				Tuple:
        _UpperCamelCase  :							str	    =	2_0
        _UpperCamelCase  :							Any	    =	model_class_name(lowerCamelCase	)

        _UpperCamelCase  :							List[str]	    =	model.init_cache(input_ids.shape[0]							,      lowerCamelCase	)
        _UpperCamelCase  :							List[Any]	    =	jnp.ones((input_ids.shape[0], max_decoder_length)							,      dtype='i4'	)

        _UpperCamelCase  :							Optional[Any]	    =	jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1	)[None, :]							,      (input_ids.shape[0], input_ids.shape[-1] - 1)	)
        _UpperCamelCase  :							Any	    =	model(
            input_ids[:, :-1]							,      attention_mask=lowerCamelCase							,      past_key_values=lowerCamelCase							,      position_ids=lowerCamelCase							,      )

        _UpperCamelCase  :							str	    =	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							,      dtype='i4'	)
        _UpperCamelCase  :							List[str]	    =	model(
            input_ids[:, -1:]							,      attention_mask=lowerCamelCase							,      past_key_values=outputs_cache.past_key_values							,      position_ids=lowerCamelCase							,      )

        _UpperCamelCase  :							Dict	    =	model(lowerCamelCase	)

        _UpperCamelCase  :							int	    =	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
        self.parent.assertTrue(diff < 1e-3							,      msg=f"Max diff is {diff}"	)

 def        __snake_case       (      self							:		Optional[int]							,      lowerCamelCase							:		Optional[Any]							,      lowerCamelCase							:		List[Any]							,      lowerCamelCase							:		int							,      lowerCamelCase							:		List[Any]	)			->				List[Any]:
        _UpperCamelCase  :							Dict	    =	2_0
        _UpperCamelCase  :							Optional[Any]	    =	model_class_name(lowerCamelCase	)

        _UpperCamelCase  :							Optional[Any]	    =	jnp.concatenate(
            [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])	)]							,      axis=-1							,      )

        _UpperCamelCase  :							Optional[Any]	    =	model.init_cache(input_ids.shape[0]							,      lowerCamelCase	)
        _UpperCamelCase  :							Tuple	    =	jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1	)[None, :]							,      (input_ids.shape[0], input_ids.shape[-1] - 1)	)

        _UpperCamelCase  :							str	    =	model(
            input_ids[:, :-1]							,      attention_mask=lowerCamelCase							,      past_key_values=lowerCamelCase							,      position_ids=lowerCamelCase							,      )
        _UpperCamelCase  :							Optional[int]	    =	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							,      dtype='i4'	)
        _UpperCamelCase  :							Tuple	    =	model(
            input_ids[:, -1:]							,      past_key_values=outputs_cache.past_key_values							,      attention_mask=lowerCamelCase							,      position_ids=lowerCamelCase							,      )

        _UpperCamelCase  :							Tuple	    =	model(lowerCamelCase							,      attention_mask=lowerCamelCase	)

        _UpperCamelCase  :							List[Any]	    =	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
        self.parent.assertTrue(diff < 1e-3							,      msg=f"Max diff is {diff}"	)







@require_flax
class  __lowerCAmelCase	(							a_,						a_,						unittest.TestCase    ):



 '''simple docstring'''
 _lowerCamelCase    :						Tuple					     =			(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
 _lowerCamelCase    :						List[str]					     =			(FlaxGPTaLMHeadModel,) if is_flax_available() else ()

 def        __snake_case       (      self							:		Dict	)			->				Union[str, Any]:
        _UpperCamelCase  :							List[str]	    =	FlaxGPTaModelTester(self	)

 def        __snake_case       (      self							:		Optional[Any]	)			->				Tuple:
        for model_class_name in self.all_model_classes:
               _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							Union[str, Any]	    =	self.model_tester.prepare_config_and_inputs()
               self.model_tester.check_use_cache_forward(lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)

 def        __snake_case       (      self							:		Union[str, Any]	)			->				List[Any]:
        for model_class_name in self.all_model_classes:
               _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							Union[str, Any]	    =	self.model_tester.prepare_config_and_inputs()
               self.model_tester.check_use_cache_forward_with_attn_mask(
                   lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)

 @slow
 def        __snake_case       (      self							:		List[str]	)			->				int:
        _UpperCamelCase  :							Any	    =	GPTaTokenizer.from_pretrained('gpt2'							,      pad_token='</s>'							,      padding_side='left'	)
        _UpperCamelCase  :							Tuple	    =	tokenizer(['Hello this is a long string', 'Hey']							,      return_tensors='np'							,      padding=lowerCamelCase							,      truncation=lowerCamelCase	)

        _UpperCamelCase  :							str	    =	FlaxGPTaLMHeadModel.from_pretrained('gpt2'	)
        _UpperCamelCase  :							Dict	    =	False
        _UpperCamelCase  :							Tuple	    =	model.config.eos_token_id

        _UpperCamelCase  :							Any	    =	jax.jit(model.generate	)

        _UpperCamelCase  :							Any	    =	jit_generate(inputs['input_ids']							,      attention_mask=inputs['attention_mask']	).sequences

        _UpperCamelCase  :							List[Any]	    =	tokenizer.batch_decode(lowerCamelCase							,      skip_special_tokens=lowerCamelCase	)

        _UpperCamelCase  :							Dict	    =	[
            'Hello this is a long string of words. I\'m going to start with the first one.\n',
            'Hey, I\'m not sure if I\'m going to be able to do',
        ]

        self.assertListEqual(lowerCamelCase							,      lowerCamelCase	)

 @is_pt_flax_cross_test
 def        __snake_case       (      self							:		int	)			->				Optional[int]:
        _UpperCamelCase    ,   _UpperCamelCase  :							Optional[Any]	    =	self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__	):
                      # prepare inputs
                      _UpperCamelCase  :							Dict	    =	self._prepare_for_class(lowerCamelCase							,      lowerCamelCase	)
                      _UpperCamelCase  :							Union[str, Any]	    =	{k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

                      # load corresponding PyTorch class
                      _UpperCamelCase  :							Tuple	    =	model_class.__name__[4:]  # Skip the "Flax" at the beginning
                      _UpperCamelCase  :							Tuple	    =	getattr(lowerCamelCase							,      lowerCamelCase	)

                      _UpperCamelCase    ,   _UpperCamelCase  :							List[str]	    =	pt_inputs['input_ids'].shape
                      _UpperCamelCase  :							Optional[Any]	    =	np.random.randint(0							,      seq_length - 1							,      size=(batch_size,)	)
                      for batch_idx, start_index in enumerate(lowerCamelCase	):
                             _UpperCamelCase  :							Any	    =	0
                             _UpperCamelCase  :							Optional[Any]	    =	1
                             _UpperCamelCase  :							Optional[int]	    =	0
                             _UpperCamelCase  :							Optional[Any]	    =	1
                      _UpperCamelCase  :							int	    =	pt_model_class(lowerCamelCase	).eval()
                      _UpperCamelCase  :							Optional[int]	    =	model_class(lowerCamelCase							,      dtype=jnp.floataa	)

                      _UpperCamelCase  :							List[str]	    =	convert_pytorch_state_dict_to_flax(pt_model.state_dict()							,      lowerCamelCase	)
                      _UpperCamelCase  :							Tuple	    =	fx_state

                      with torch.no_grad():
                             _UpperCamelCase  :							Optional[int]	    =	pt_model(**lowerCamelCase	).to_tuple()

                      _UpperCamelCase  :							Any	    =	fx_model(**lowerCamelCase	).to_tuple()
                      self.assertEqual(len(lowerCamelCase	)							,      len(lowerCamelCase	)							,      'Output lengths differ between Flax and PyTorch'	)
                      for fx_output, pt_output in zip(lowerCamelCase							,      lowerCamelCase	):
                             self.assert_almost_equals(fx_output[:, -1]							,      pt_output[:, -1].numpy()							,      4e-2	)

                      with tempfile.TemporaryDirectory() as tmpdirname:
                             pt_model.save_pretrained(lowerCamelCase	)
                             _UpperCamelCase  :							Optional[Any]	    =	model_class.from_pretrained(lowerCamelCase							,      from_pt=lowerCamelCase	)

                      _UpperCamelCase  :							List[Any]	    =	fx_model_loaded(**lowerCamelCase	).to_tuple()
                      self.assertEqual(
                          len(lowerCamelCase	)							,      len(lowerCamelCase	)							,      'Output lengths differ between Flax and PyTorch'	)
                      for fx_output_loaded, pt_output in zip(lowerCamelCase							,      lowerCamelCase	):
                             self.assert_almost_equals(fx_output_loaded[:, -1]							,      pt_output[:, -1].numpy()							,      4e-2	)

 @is_pt_flax_cross_test
 def        __snake_case       (      self							:		Optional[Any]	)			->				int:
        _UpperCamelCase    ,   _UpperCamelCase  :							Any	    =	self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__	):
                      # prepare inputs
                      _UpperCamelCase  :							Optional[int]	    =	self._prepare_for_class(lowerCamelCase							,      lowerCamelCase	)
                      _UpperCamelCase  :							int	    =	{k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

                      # load corresponding PyTorch class
                      _UpperCamelCase  :							str	    =	model_class.__name__[4:]  # Skip the "Flax" at the beginning
                      _UpperCamelCase  :							Union[str, Any]	    =	getattr(lowerCamelCase							,      lowerCamelCase	)

                      _UpperCamelCase  :							Optional[Any]	    =	pt_model_class(lowerCamelCase	).eval()
                      _UpperCamelCase  :							str	    =	model_class(lowerCamelCase							,      dtype=jnp.floataa	)

                      _UpperCamelCase  :							Tuple	    =	load_flax_weights_in_pytorch_model(lowerCamelCase							,      fx_model.params	)
                      _UpperCamelCase    ,   _UpperCamelCase  :							List[Any]	    =	pt_inputs['input_ids'].shape
                      _UpperCamelCase  :							List[Any]	    =	np.random.randint(0							,      seq_length - 1							,      size=(batch_size,)	)
                      for batch_idx, start_index in enumerate(lowerCamelCase	):
                             _UpperCamelCase  :							Optional[int]	    =	0
                             _UpperCamelCase  :							Union[str, Any]	    =	1
                             _UpperCamelCase  :							Dict	    =	0
                             _UpperCamelCase  :							Union[str, Any]	    =	1

                      # make sure weights are tied in PyTorch
                      pt_model.tie_weights()

                      with torch.no_grad():
                             _UpperCamelCase  :							Optional[int]	    =	pt_model(**lowerCamelCase	).to_tuple()

                      _UpperCamelCase  :							Any	    =	fx_model(**lowerCamelCase	).to_tuple()
                      self.assertEqual(len(lowerCamelCase	)							,      len(lowerCamelCase	)							,      'Output lengths differ between Flax and PyTorch'	)
                      for fx_output, pt_output in zip(lowerCamelCase							,      lowerCamelCase	):
                             self.assert_almost_equals(fx_output[:, -1]							,      pt_output[:, -1].numpy()							,      4e-2	)

                      with tempfile.TemporaryDirectory() as tmpdirname:
                             fx_model.save_pretrained(lowerCamelCase	)
                             _UpperCamelCase  :							Dict	    =	pt_model_class.from_pretrained(lowerCamelCase							,      from_flax=lowerCamelCase	)

                      with torch.no_grad():
                             _UpperCamelCase  :							Tuple	    =	pt_model_loaded(**lowerCamelCase	).to_tuple()

                      self.assertEqual(
                          len(lowerCamelCase	)							,      len(lowerCamelCase	)							,      'Output lengths differ between Flax and PyTorch'	)
                      for fx_output, pt_output in zip(lowerCamelCase							,      lowerCamelCase	):
                             self.assert_almost_equals(fx_output[:, -1]							,      pt_output[:, -1].numpy()							,      4e-2	)

 @slow
 def        __snake_case       (      self							:		Union[str, Any]	)			->				Dict:
        for model_class_name in self.all_model_classes:
               _UpperCamelCase  :							List[Any]	    =	model_class_name.from_pretrained('gpt2'							,      from_pt=lowerCamelCase	)
               _UpperCamelCase  :							Tuple	    =	model(np.ones((1, 1)	)	)
               self.assertIsNotNone(lowerCamelCase	)





'''simple docstring'''







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






					def __init__(  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE=2							,				__SCREAMING_SNAKE_CASE=True							,				__SCREAMING_SNAKE_CASE=True							,				__SCREAMING_SNAKE_CASE=10							,				__SCREAMING_SNAKE_CASE=10							,				__SCREAMING_SNAKE_CASE=8							,				__SCREAMING_SNAKE_CASE=[2, 4, 4]							,				__SCREAMING_SNAKE_CASE=3							,				__SCREAMING_SNAKE_CASE=7_68							,				__SCREAMING_SNAKE_CASE=5							,				__SCREAMING_SNAKE_CASE=4							,				__SCREAMING_SNAKE_CASE=37							,				__SCREAMING_SNAKE_CASE="gelu_fast"							,				__SCREAMING_SNAKE_CASE=0.0							,				__SCREAMING_SNAKE_CASE=0.0							,				__SCREAMING_SNAKE_CASE=0.02							,				__SCREAMING_SNAKE_CASE=1E-0_6							,				__SCREAMING_SNAKE_CASE=True							,				__SCREAMING_SNAKE_CASE=None							,				):


											"""simple docstring"""





											lowercase_      :	Any						=			parent
											lowercase_      :	Union[str, Any]						=			batch_size
											lowercase_      :	Union[str, Any]						=			is_training
											lowercase_      :	Tuple						=			use_labels
											lowercase_      :	Any						=			num_labels
											lowercase_      :	str						=			image_size
											lowercase_      :	Union[str, Any]						=			num_frames
											lowercase_      :	Any						=			tubelet_size
											lowercase_      :	Tuple						=			num_channels
											lowercase_      :	Union[str, Any]						=			hidden_size
											lowercase_      :	Dict						=			num_hidden_layers
											lowercase_      :	Optional[int]						=			num_attention_heads
											lowercase_      :	str						=			intermediate_size
											lowercase_      :	Dict						=			hidden_act
											lowercase_      :	Dict						=			hidden_dropout_prob
											lowercase_      :	Optional[Any]						=			attention_probs_dropout_prob
											lowercase_      :	Dict						=			initializer_range
											lowercase_      :	Optional[int]						=			layer_norm_eps
											lowercase_      :	str						=			qkv_bias
											lowercase_      :	Dict						=			scope

											lowercase_      :	Dict						=			(
											    (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       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	Tuple						=			floats_tensor(
											    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]						)

											lowercase_      :	Optional[int]						=			None
											if self.use_labels:
																	lowercase_      :	Tuple						=			ids_tensor([self.batch_size]							,				self.num_labels						)

											lowercase_      :	Union[str, Any]						=			self.get_config()

											return config, pixel_values, labels






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	Optional[int]						=			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							,				)
											lowercase_      :	Optional[Any]						=			self.num_labels
											return config






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):


											"""simple docstring"""





											lowercase_      :	str						=			VivitModel(config=__SCREAMING_SNAKE_CASE						)
											model.to(__SCREAMING_SNAKE_CASE						)
											model.eval()
											lowercase_      :	Union[str, Any]						=			model(__SCREAMING_SNAKE_CASE						)
											self.parent.assertEqual(result.last_hidden_state.shape							,				(self.batch_size, self.seq_length, self.hidden_size)						)






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):


											"""simple docstring"""





											lowercase_      :	List[Any]						=			VivitForVideoClassification(__SCREAMING_SNAKE_CASE						)
											model.to(__SCREAMING_SNAKE_CASE						)
											model.eval()

											lowercase_      :	Union[str, Any]						=			model(__SCREAMING_SNAKE_CASE						)

											# verify the logits shape
											lowercase_      :	Union[str, Any]						=			torch.Size((self.batch_size, self.num_labels)						)
											self.parent.assertEqual(result.logits.shape							,				__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	int						=			self.prepare_config_and_inputs()
											lowercase_		,					lowercase_		,					lowercase_      :	Dict						=			config_and_inputs
											lowercase_      :	List[str]						=			{'''pixel_values''': pixel_values}
											return config, inputs_dict







@require_torch
class 					lowerCAmelCase__  (    lowerCamelCase_       ,		unittest.TestCase	):
					lowerCAmelCase_	  =	(VivitModel, VivitForVideoClassification) if is_torch_available() else ()

					lowerCAmelCase_	  =	False
					lowerCAmelCase_	  =	False
					lowerCAmelCase_	  =	False
					lowerCAmelCase_	  =	False






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	int						=			VivitModelTester(self						)
											lowercase_      :	Optional[int]						=			ConfigTester(self							,				config_class=__SCREAMING_SNAKE_CASE							,				has_text_modality=__SCREAMING_SNAKE_CASE							,				hidden_size=37						)






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE=False						):


											"""simple docstring"""





											lowercase_      :	str						=			copy.deepcopy(__SCREAMING_SNAKE_CASE						)

											if return_labels:
																	if model_class in get_values(__SCREAMING_SNAKE_CASE						):
																							lowercase_      :	List[str]						=			torch.zeros(
																							    self.model_tester.batch_size							,				dtype=torch.long							,				device=__SCREAMING_SNAKE_CASE						)

											return inputs_dict






					def       _snake_case       (  self						):


											"""simple docstring"""





											self.config_tester.run_common_tests()






					@unittest.skip(reason='''Vivit does not use inputs_embeds'''						)
					def       _snake_case       (  self						):


											"""simple docstring"""





											pass






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_		,					lowercase_      :	str						=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowercase_      :	Optional[Any]						=			model_class(__SCREAMING_SNAKE_CASE						)
																	self.assertIsInstance(model.get_input_embeddings()							,				(nn.Module)						)
																	lowercase_      :	int						=			model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(__SCREAMING_SNAKE_CASE							,				nn.Linear						)						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_		,					lowercase_      :	str						=			self.model_tester.prepare_config_and_inputs_for_common()

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

																	lowercase_      :	str						=			['''pixel_values''', '''head_mask''']
																	self.assertListEqual(arg_names[:2]							,				__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	Union[str, Any]						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	int						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_video_classification(*__SCREAMING_SNAKE_CASE						)






					@slow
					def       _snake_case       (  self						):


											"""simple docstring"""





											for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase_      :	Union[str, Any]						=			VivitModel.from_pretrained(__SCREAMING_SNAKE_CASE						)
																	self.assertIsNotNone(__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_		,					lowercase_      :	Union[str, Any]						=			self.model_tester.prepare_config_and_inputs_for_common()
											lowercase_      :	Union[str, Any]						=			True

											for model_class in self.all_model_classes:
																	lowercase_      :	int						=			self.model_tester.seq_length

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

																	# check that output_attentions also work using config
																	del inputs_dict["output_attentions"]
																	lowercase_      :	List[Any]						=			True
																	lowercase_      :	Union[str, Any]						=			model_class(__SCREAMING_SNAKE_CASE						)
																	model.to(__SCREAMING_SNAKE_CASE						)
																	model.eval()
																	with torch.no_grad():
																							lowercase_      :	str						=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)						)
																	lowercase_      :	Dict						=			outputs.attentions
																	self.assertEqual(len(__SCREAMING_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]							,				)
																	lowercase_      :	Dict						=			len(__SCREAMING_SNAKE_CASE						)

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

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

																	lowercase_      :	List[str]						=			outputs.attentions

																	self.assertEqual(len(__SCREAMING_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						):


											"""simple docstring"""





											def check_hidden_states_output(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):
																	lowercase_      :	List[str]						=			model_class(__SCREAMING_SNAKE_CASE						)
																	model.to(__SCREAMING_SNAKE_CASE						)
																	model.eval()

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

																	lowercase_      :	List[str]						=			outputs.hidden_states
																	lowercase_      :	Any						=			self.model_tester.num_hidden_layers + 1
																	self.assertEqual(len(__SCREAMING_SNAKE_CASE						)							,				__SCREAMING_SNAKE_CASE						)

																	lowercase_      :	List[Any]						=			self.model_tester.seq_length

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

											lowercase_		,					lowercase_      :	List[str]						=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowercase_      :	int						=			True
																	check_hidden_states_output(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)

																	# check that output_hidden_states also work using config
																	del inputs_dict["output_hidden_states"]
																	lowercase_      :	List[Any]						=			True

																	check_hidden_states_output(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)






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


											"""simple docstring"""





											pass







def    snake_case_						(							):



						"""simple docstring"""





						lowercase_      :	List[str]						=			hf_hub_download(
						    repo_id='''hf-internal-testing/spaghetti-video''' , filename='''eating_spaghetti_32_frames.npy''' , repo_type='''dataset'''				)
						lowercase_      :	List[str]						=			np.load(__SCREAMING_SNAKE_CASE				)
						return list(__SCREAMING_SNAKE_CASE				)






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






					@cached_property
					def       _snake_case       (  self						):


											"""simple docstring"""





											return VivitImageProcessor() if is_vision_available() else None






					@slow
					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	List[str]						=			VivitForVideoClassification.from_pretrained('''google/vivit-b-16x2-kinetics400'''						).to(__SCREAMING_SNAKE_CASE						)

											lowercase_      :	Optional[int]						=			self.default_image_processor
											lowercase_      :	Union[str, Any]						=			prepare_video()
											lowercase_      :	str						=			image_processor(__SCREAMING_SNAKE_CASE							,				return_tensors='''pt'''						).to(__SCREAMING_SNAKE_CASE						)

											# forward pass
											with torch.no_grad():
																	lowercase_      :	Any						=			model(**__SCREAMING_SNAKE_CASE						)

											# verify the logits
											lowercase_      :	Any						=			torch.Size((1, 4_00)						)
											self.assertEqual(outputs.logits.shape							,				__SCREAMING_SNAKE_CASE						)

											# taken from original model
											lowercase_      :	Any						=			torch.tensor([-1.0_543, 2.0_764, -0.2_104, 0.4_439, -0.9_658]						).to(__SCREAMING_SNAKE_CASE						)

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



