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				,  lowerCAmelCase_				,  lowerCAmelCase_=2				,  lowerCAmelCase_=True				,  lowerCAmelCase_=True				,  lowerCAmelCase_=10				,  lowerCAmelCase_=10				,  lowerCAmelCase_=8				,  lowerCAmelCase_=[2, 4, 4]				,  lowerCAmelCase_=3				,  lowerCAmelCase_=7_68				,  lowerCAmelCase_=5				,  lowerCAmelCase_=4				,  lowerCAmelCase_=37				,  lowerCAmelCase_="gelu_fast"				,  lowerCAmelCase_=0.0				,  lowerCAmelCase_=0.0				,  lowerCAmelCase_=0.02				,  lowerCAmelCase_=1e-06				,  lowerCAmelCase_=True				,  lowerCAmelCase_=None				,  )						->Dict:
        a	:     List[str]						      =    parent
        a	:     int						      =    batch_size
        a	:     Dict						      =    is_training
        a	:     str						      =    use_labels
        a	:     Dict						      =    num_labels
        a	:     Any						      =    image_size
        a	:     Union[str, Any]						      =    num_frames
        a	:     int						      =    tubelet_size
        a	:     str						      =    num_channels
        a	:     Optional[int]						      =    hidden_size
        a	:     Union[str, Any]						      =    num_hidden_layers
        a	:     List[Any]						      =    num_attention_heads
        a	:     Optional[Any]						      =    intermediate_size
        a	:     Tuple						      =    hidden_act
        a	:     Dict						      =    hidden_dropout_prob
        a	:     Any						      =    attention_probs_dropout_prob
        a	:     List[str]						      =    initializer_range
        a	:     Optional[Any]						      =    layer_norm_eps
        a	:     Union[str, Any]						      =    qkv_bias
        a	:     Tuple						      =    scope

        a	:     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 _lowercase(     self     )						->Optional[int]:
        a	:     Union[str, Any]						      =    floats_tensor(
            [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]     )

        a	:     Union[str, Any]						      =    None
        if self.use_labels:
             a	:     Dict						      =    ids_tensor([self.batch_size]				,  self.num_labels     )

        a	:     Optional[Any]						      =    self.get_config()

        return config, pixel_values, labels
   def _lowercase(     self     )						->List[str]:
        a	:     List[str]						      =    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	:     int						      =    self.num_labels
        return config
   def _lowercase(     self				,  lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_     )						->Optional[int]:
        a	:     Union[str, Any]						      =    VivitModel(config=lowerCAmelCase_     )
        model.to(lowerCAmelCase_     )
        model.eval()
        a	:     List[Any]						      =    model(lowerCAmelCase_     )
        self.parent.assertEqual(result.last_hidden_state.shape				,  (self.batch_size, self.seq_length, self.hidden_size)     )
   def _lowercase(     self				,  lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_     )						->Any:
        a	:     List[Any]						      =    VivitForVideoClassification(lowerCAmelCase_     )
        model.to(lowerCAmelCase_     )
        model.eval()

        a	:     Optional[int]						      =    model(lowerCAmelCase_     )

        # verify the logits shape
        a	:     Optional[int]						      =    torch.Size((self.batch_size, self.num_labels)     )
        self.parent.assertEqual(result.logits.shape				,  lowerCAmelCase_     )


   def _lowercase(     self     )						->List[str]:
        a	:     Tuple						      =    self.prepare_config_and_inputs()
        a      ,				a      ,				a	:     Dict						      =    config_and_inputs
        a	:     str						      =    {'''pixel_values''': pixel_values}
        return config, inputs_dict






@require_torch
class 				_SCREAMING_SNAKE_CASE( snake_case__							,   unittest.TestCase						):
   UpperCamelCase_					:    Optional[int]			      = (VivitModel, VivitForVideoClassification) if is_torch_available() else ()

   UpperCamelCase_					:    int			      = False
   UpperCamelCase_					:    Any			      = False
   UpperCamelCase_					:    Tuple			      = False
   UpperCamelCase_					:    Dict			      = False
   def _lowercase(     self     )						->Optional[int]:
        a	:     Optional[int]						      =    VivitModelTester(self     )
        a	:     int						      =    ConfigTester(self				,  config_class=lowerCAmelCase_				,  has_text_modality=lowerCAmelCase_				,  hidden_size=37     )
   def _lowercase(     self				,  lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_=False     )						->int:
        a	:     int						      =    copy.deepcopy(lowerCAmelCase_     )

        if return_labels:
             if model_class in get_values(lowerCAmelCase_     ):
                  a	:     Any						      =    torch.zeros(
                      self.model_tester.batch_size				,  dtype=torch.long				,  device=lowerCAmelCase_     )

        return inputs_dict
   def _lowercase(     self     )						->Dict:
        self.config_tester.run_common_tests()
   @unittest.skip(reason='''Vivit does not use inputs_embeds'''     )
   def _lowercase(     self     )						->List[Any]:
        pass
   def _lowercase(     self     )						->int:
        a      ,				a	:     Union[str, Any]						      =    self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             a	:     Optional[int]						      =    model_class(lowerCAmelCase_     )
             self.assertIsInstance(model.get_input_embeddings()				,  (nn.Module)     )
             a	:     Union[str, Any]						      =    model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(lowerCAmelCase_				,  nn.Linear     )     )
   def _lowercase(     self     )						->Optional[Any]:
        a      ,				a	:     Tuple						      =    self.model_tester.prepare_config_and_inputs_for_common()

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

             a	:     Dict						      =    ['''pixel_values''', '''head_mask''']
             self.assertListEqual(arg_names[:2]				,  lowerCAmelCase_     )
   def _lowercase(     self     )						->int:
        a	:     Tuple						      =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase_     )
   def _lowercase(     self     )						->int:
        a	:     Union[str, Any]						      =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_video_classification(*lowerCAmelCase_     )
   @slow
   def _lowercase(     self     )						->Any:
        for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             a	:     Union[str, Any]						      =    VivitModel.from_pretrained(lowerCAmelCase_     )
             self.assertIsNotNone(lowerCAmelCase_     )
   def _lowercase(     self     )						->List[Any]:
        a      ,				a	:     List[Any]						      =    self.model_tester.prepare_config_and_inputs_for_common()
        a	:     Dict						      =    True

        for model_class in self.all_model_classes:
             a	:     Union[str, Any]						      =    self.model_tester.seq_length

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

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

             self.assertListEqual(
                 list(attentions[0].shape[-3:]     )				,  [self.model_tester.num_attention_heads, seq_len, seq_len]				,  )
             a	:     List[str]						      =    len(lowerCAmelCase_     )

             # Check attention is always last and order is fine
             a	:     Tuple						      =    True
             a	:     Optional[Any]						      =    True
             a	:     Optional[Any]						      =    model_class(lowerCAmelCase_     )
             model.to(lowerCAmelCase_     )
             model.eval()
             with torch.no_grad():
                  a	:     Optional[Any]						      =    model(**self._prepare_for_class(lowerCAmelCase_				,  lowerCAmelCase_     )     )

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

             a	:     List[Any]						      =    outputs.attentions

             self.assertEqual(len(lowerCAmelCase_     )				,  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 _lowercase(     self     )						->str:
        def check_hidden_states_output(lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_     ):
             a	:     Any						      =    model_class(lowerCAmelCase_     )
             model.to(lowerCAmelCase_     )
             model.eval()

             with torch.no_grad():
                  a	:     int						      =    model(**self._prepare_for_class(lowerCAmelCase_				,  lowerCAmelCase_     )     )

             a	:     Tuple						      =    outputs.hidden_states
             a	:     int						      =    self.model_tester.num_hidden_layers + 1
             self.assertEqual(len(lowerCAmelCase_     )				,  lowerCAmelCase_     )

             a	:     Union[str, Any]						      =    self.model_tester.seq_length

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

        a      ,				a	:     int						      =    self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             a	:     Optional[int]						      =    True
             check_hidden_states_output(lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_     )

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

             check_hidden_states_output(lowerCAmelCase_				,  lowerCAmelCase_				,  lowerCAmelCase_     )


   @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''     )
   def _lowercase(     self     )						->Union[str, Any]:
        pass







def 						UpperCAmelCase							(		)						->						str:
     a	:     Optional[Any]						      =    hf_hub_download(
         repo_id='''hf-internal-testing/spaghetti-video'''    ,   filename='''eating_spaghetti_32_frames.npy'''    ,   repo_type='''dataset'''							)
     a	:     Union[str, Any]						      =    np.load(lowerCAmelCase_							)
     return list(lowerCAmelCase_							)






@require_torch
@require_vision
class 				_SCREAMING_SNAKE_CASE( unittest.TestCase						):
   @cached_property
   def _lowercase(     self     )						->List[Any]:
        return VivitImageProcessor() if is_vision_available() else None


   @slow
   def _lowercase(     self     )						->List[str]:
        a	:     Tuple						      =    VivitForVideoClassification.from_pretrained('''google/vivit-b-16x2-kinetics400'''     ).to(lowerCAmelCase_     )

        a	:     int						      =    self.default_image_processor
        a	:     Dict						      =    prepare_video()
        a	:     str						      =    image_processor(lowerCAmelCase_				,  return_tensors='''pt'''     ).to(lowerCAmelCase_     )

        # forward pass
        with torch.no_grad():
             a	:     Dict						      =    model(**lowerCAmelCase_     )

        # verify the logits
        a	:     int						      =    torch.Size((1, 4_00)     )
        self.assertEqual(outputs.logits.shape				,  lowerCAmelCase_     )

        # taken from original model
        a	:     int						      =    torch.tensor([-1.05_43, 2.07_64, -0.21_04, 0.44_39, -0.96_58]     ).to(lowerCAmelCase_     )

        self.assertTrue(torch.allclose(outputs.logits[0, :5]				,  lowerCAmelCase_				,  atol=1e-4     )     )
from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase				 =   logging.get_logger(__name__)

UpperCAmelCase				 =   {
    """facebook/opt-125m""": """https://huggingface.co/facebook/opt-125m/blob/main/config.json""",
    """facebook/opt-350m""": """https://huggingface.co/facebook/opt-350m/blob/main/config.json""",
    """facebook/opt-1.3b""": """https://huggingface.co/facebook/opt-1.3b/blob/main/config.json""",
    """facebook/opt-2.7b""": """https://huggingface.co/facebook/opt-2.7b/blob/main/config.json""",
    """facebook/opt-6.7b""": """https://huggingface.co/facebook/opt-6.7b/blob/main/config.json""",
    """facebook/opt-13b""": """https://huggingface.co/facebook/opt-13b/blob/main/config.json""",
}






class 				_SCREAMING_SNAKE_CASE( snake_case__						):
   UpperCamelCase_					:    List[str]			      = 'opt'
   UpperCamelCase_					:    List[str]			      = ['past_key_values']
   def __init__(     self				,  lowerCAmelCase_=5_02_72				,  lowerCAmelCase_=7_68				,  lowerCAmelCase_=12				,  lowerCAmelCase_=30_72				,  lowerCAmelCase_=20_48				,  lowerCAmelCase_=True				,  lowerCAmelCase_=False				,  lowerCAmelCase_=None				,  lowerCAmelCase_=0.1				,  lowerCAmelCase_=0.0				,  lowerCAmelCase_=12				,  lowerCAmelCase_="relu"				,  lowerCAmelCase_=0.0				,  lowerCAmelCase_=0.02				,  lowerCAmelCase_=True				,  lowerCAmelCase_=1				,  lowerCAmelCase_=2				,  lowerCAmelCase_=2				,  lowerCAmelCase_=True				,  lowerCAmelCase_=True				,  **lowerCAmelCase_				,  )						->Any:
        super().__init__(
            pad_token_id=lowerCAmelCase_				,  bos_token_id=lowerCAmelCase_				,  eos_token_id=lowerCAmelCase_				,  **lowerCAmelCase_				,  )
        a	:     Dict						      =    vocab_size
        a	:     int						      =    max_position_embeddings
        a	:     List[Any]						      =    num_attention_heads
        a	:     Union[str, Any]						      =    word_embed_proj_dim if word_embed_proj_dim is not None else hidden_size
        a	:     Any						      =    ffn_dim
        a	:     Optional[int]						      =    hidden_size
        a	:     Optional[Any]						      =    num_hidden_layers
        a	:     Optional[int]						      =    dropout
        a	:     Optional[Any]						      =    attention_dropout
        a	:     Union[str, Any]						      =    activation_function
        a	:     Any						      =    init_std
        a	:     List[Any]						      =    layerdrop
        a	:     str						      =    use_cache
        a	:     Any						      =    do_layer_norm_before
        # We keep these variables at `True` for backward compatibility.
        a	:     List[Any]						      =    enable_bias
        a	:     List[str]						      =    layer_norm_elementwise_affine

        # Note that the only purpose of `_remove_final_layer_norm` is to keep backward compatibility
        # with checkpoints that have been fine-tuned before transformers v4.20.1
        # see https://github.com/facebookresearch/metaseq/pull/164
        a	:     Dict						      =    _remove_final_layer_norm