


'''simple docstring'''






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 ,			__magic_name__ ,			__magic_name__=13 ,			__magic_name__=10 ,			__magic_name__=3 ,			__magic_name__=2 ,			__magic_name__=2 ,			__magic_name__=2 ,			__magic_name__=True ,			__magic_name__=True ,			__magic_name__=32 ,			__magic_name__=5 ,			__magic_name__=4 ,			__magic_name__=37 ,			__magic_name__="gelu" ,			__magic_name__=0.1 ,			__magic_name__=0.1 ,			__magic_name__=10 ,			__magic_name__=0.0_2 ,			__magic_name__=0.9 ,			__magic_name__=None ,			)   ->							Optional[int]:
      _a      =							parent
      _a      =							batch_size
      _a      =							image_size
      _a      =							num_channels
      _a      =							patch_size
      _a      =							tubelet_size
      _a      =							num_frames
      _a      =							is_training
      _a      =							use_labels
      _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      =							type_sequence_label_size
      _a      =							initializer_range
      _a      =							mask_ratio
      _a      =							scope

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

      # use this variable to define bool_masked_pos
      _a      =							int(mask_ratio * self.seq_length  )
 def 							__UpperCAmelCase						(	self  )   ->							Optional[int]:
      _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.type_sequence_label_size  )

      _a      =							self.get_config()

      return config, pixel_values, labels
 def 							__UpperCAmelCase						(	self  )   ->							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=__magic_name__ ,			initializer_range=self.initializer_range ,			)
 def 							__UpperCAmelCase						(	self ,			__magic_name__ ,			__magic_name__ ,			__magic_name__  )   ->							Union[str, Any]:
      _a      =							VideoMAEModel(config=__magic_name__  )
      model.to(__magic_name__  )
      model.eval()
      _a      =							model(__magic_name__  )
      self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, self.seq_length, self.hidden_size)  )
 def 							__UpperCAmelCase						(	self ,			__magic_name__ ,			__magic_name__ ,			__magic_name__  )   ->							str:
      _a      =							VideoMAEForPreTraining(__magic_name__  )
      model.to(__magic_name__  )
      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
      _a      =							torch.ones((self.num_masks,)  )
      _a      =							torch.cat([mask, torch.zeros(self.seq_length - mask.size(0  )  )]  )
      _a      =							mask.expand(self.batch_size ,			-1  ).bool()

      _a      =							model(__magic_name__ ,			__magic_name__  )
      # model only returns predictions for masked patches
      _a      =							mask.sum().item()
      _a      =							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  )   ->							Optional[int]:
      _a      =							self.prepare_config_and_inputs()
      _a ,				_a ,				_a      =							config_and_inputs
      _a      =							{'pixel_values': pixel_values}
      return config, inputs_dict





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

 _lowerCAmelCase	   =		False
 _lowerCAmelCase	   =		False
 _lowerCAmelCase	   =		False
 _lowerCAmelCase	   =		False
 def 							__UpperCAmelCase						(	self  )   ->							str:
      _a      =							VideoMAEModelTester(self  )
      _a      =							ConfigTester(self ,			config_class=__magic_name__ ,			has_text_modality=__magic_name__ ,			hidden_size=37  )
 def 							__UpperCAmelCase						(	self ,			__magic_name__ ,			__magic_name__ ,			__magic_name__=False  )   ->							List[Any]:
      _a      =							copy.deepcopy(__magic_name__  )

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

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

      return inputs_dict
 def 							__UpperCAmelCase						(	self  )   ->							Optional[Any]:
      self.config_tester.run_common_tests()
 @unittest.skip(reason='VideoMAE does not use inputs_embeds'  )
 def 							__UpperCAmelCase						(	self  )   ->							Optional[int]:
      pass
 def 							__UpperCAmelCase						(	self  )   ->							List[str]:
      _a ,				_a      =							self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           _a      =							model_class(__magic_name__  )
           self.assertIsInstance(model.get_input_embeddings() ,			(nn.Module)  )
           _a      =							model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(__magic_name__ ,			nn.Linear  )  )
 def 							__UpperCAmelCase						(	self  )   ->							Any:
      _a ,				_a      =							self.model_tester.prepare_config_and_inputs_for_common()

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

           _a      =							['pixel_values']
           self.assertListEqual(arg_names[:1] ,			__magic_name__  )
 def 							__UpperCAmelCase						(	self  )   ->							Optional[int]:
      _a      =							self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__magic_name__  )
 def 							__UpperCAmelCase						(	self  )   ->							Any:
      _a      =							self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_pretraining(*__magic_name__  )
 @slow
 def 							__UpperCAmelCase						(	self  )   ->							str:
      for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _a      =							VideoMAEModel.from_pretrained(__magic_name__  )
           self.assertIsNotNone(__magic_name__  )
 def 							__UpperCAmelCase						(	self  )   ->							Union[str, Any]:
      if not self.has_attentions:
           pass

      else:
           _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 - self.model_tester.num_masks
                _a      =							(
                    num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
                )

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

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

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

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

                _a      =							outputs.attentions

                self.assertEqual(len(__magic_name__  ) ,			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:
      def check_hidden_states_output(__magic_name__ ,			__magic_name__ ,			__magic_name__  ):
           _a      =							model_class(__magic_name__  )
           model.to(__magic_name__  )
           model.eval()

           with torch.no_grad():
                _a      =							model(**self._prepare_for_class(__magic_name__ ,			__magic_name__  )  )

           _a      =							outputs.hidden_states
           _a      =							self.model_tester.num_hidden_layers + 1
           self.assertEqual(len(__magic_name__  ) ,			__magic_name__  )

           _a      =							self.model_tester.seq_length - self.model_tester.num_masks
           _a      =							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] ,			)

      _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(__magic_name__ ,			__magic_name__ ,			__magic_name__  )

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

           check_hidden_states_output(__magic_name__ ,			__magic_name__ ,			__magic_name__  )




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






def 		_A			()				->					Dict:




     '''simple docstring'''



     _a      =							hf_hub_download(
         repo_id='hf-internal-testing/spaghetti-video'			,  filename='eating_spaghetti.npy'			,  repo_type='dataset'						)
     _a      =							np.load(lowerCAmelCase__						)
     return list(lowerCAmelCase__						)






@require_torch
@require_vision
class 	a       (       unittest.TestCase    ):
 @cached_property
 def 							__UpperCAmelCase						(	self  )   ->							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  )   ->							Tuple:
      _a      =							VideoMAEForVideoClassification.from_pretrained('MCG-NJU/videomae-base-finetuned-kinetics'  ).to(
          __magic_name__  )

      _a      =							self.default_image_processor
      _a      =							prepare_video()
      _a      =							image_processor(__magic_name__ ,			return_tensors='pt'  ).to(__magic_name__  )

      # forward pass
      with torch.no_grad():
           _a      =							model(**__magic_name__  )

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

      _a      =							torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]  ).to(__magic_name__  )

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




 @slow
 def 							__UpperCAmelCase						(	self  )   ->							Any:
      _a      =							VideoMAEForPreTraining.from_pretrained('MCG-NJU/videomae-base-short'  ).to(__magic_name__  )

      _a      =							self.default_image_processor
      _a      =							prepare_video()
      _a      =							image_processor(__magic_name__ ,			return_tensors='pt'  ).to(__magic_name__  )

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

      # forward pass
      with torch.no_grad():
           _a      =							model(**__magic_name__  )

      # verify the logits
      _a      =							torch.Size([1, 14_08, 15_36]  )
      _a      =							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=__magic_name__  )
      self.assertEqual(outputs.logits.shape ,			__magic_name__  )
      self.assertTrue(torch.allclose(outputs.logits[0, :3, :3] ,			__magic_name__ ,			atol=1e-4  )  )

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

      # verify the loss (`config.norm_pix_loss` = `False`)
      _a      =							VideoMAEForPreTraining.from_pretrained('MCG-NJU/videomae-base-short' ,			norm_pix_loss=__magic_name__  ).to(
          __magic_name__  )

      with torch.no_grad():
           _a      =							model(**__magic_name__  )

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





'''simple docstring'''






from __future__ import annotations

from sys import maxsize
from typing import Generic, TypeVar

a_   :      Optional[Any]				    =      TypeVar("T")






def 		_A			(lowerCAmelCase__					:int						)				->					int:




     '''simple docstring'''



     return (position - 1) // 2






def 		_A			(lowerCAmelCase__					:int						)				->					int:




     '''simple docstring'''



     return (2 * position) + 1






def 		_A			(lowerCAmelCase__					:int						)				->					int:




     '''simple docstring'''



     return (2 * position) + 2




class 	a       (       Generic[T]    ):
 def __init__(	self  )   ->							None:
      _a      =							[]
      _a      =							{}
      _a      =							0
 def __len__(	self  )   ->							int:
      return self.elements
 def __repr__(	self  )   ->							str:
      return str(self.heap  )
 def 							__UpperCAmelCase						(	self  )   ->							bool:
      # Check if the priority queue is empty
      return self.elements == 0
 def 							__UpperCAmelCase						(	self ,			__magic_name__ ,			__magic_name__  )   ->							None:
      # Add an element with given priority to the queue
      self.heap.append((elem, weight)  )
      _a      =							self.elements
      self.elements += 1
      self._bubble_up(__magic_name__  )
 def 							__UpperCAmelCase						(	self  )   ->							T:
      # Remove and return the element with lowest weight (highest priority)
      if self.elements > 1:
           self._swap_nodes(0 ,			self.elements - 1  )
      _a ,				_a      =							self.heap.pop()
      del self.position_map[elem]
      self.elements -= 1
      if self.elements > 0:
           _a ,				_a      =							self.heap[0]
           self._bubble_down(__magic_name__  )
      return elem
 def 							__UpperCAmelCase						(	self ,			__magic_name__ ,			__magic_name__  )   ->							None:
      # Update the weight of the given key
      _a      =							self.position_map[elem]
      _a      =							(elem, weight)
      if position > 0:
           _a      =							get_parent_position(__magic_name__  )
           _a ,				_a      =							self.heap[parent_position]
           if parent_weight > weight:
                self._bubble_up(__magic_name__  )
           else:
                self._bubble_down(__magic_name__  )
      else:
           self._bubble_down(__magic_name__  )
 def 							__UpperCAmelCase						(	self ,			__magic_name__  )   ->							None:
      # Place a node at the proper position (upward movement) [to be used internally
      # only]
      _a      =							self.position_map[elem]
      if curr_pos == 0:
           return None
      _a      =							get_parent_position(__magic_name__  )
      _a ,				_a      =							self.heap[curr_pos]
      _a ,				_a      =							self.heap[parent_position]
      if parent_weight > weight:
           self._swap_nodes(__magic_name__ ,			__magic_name__  )
           return self._bubble_up(__magic_name__  )
      return None
 def 							__UpperCAmelCase						(	self ,			__magic_name__  )   ->							None:
      # Place a node at the proper position (downward movement) [to be used
      # internally only]
      _a      =							self.position_map[elem]
      _a ,				_a      =							self.heap[curr_pos]
      _a      =							get_child_left_position(__magic_name__  )
      _a      =							get_child_right_position(__magic_name__  )
      if child_left_position < self.elements and child_right_position < self.elements:
           _a ,				_a      =							self.heap[child_left_position]
           _a ,				_a      =							self.heap[child_right_position]
           if child_right_weight < child_left_weight and child_right_weight < weight:
                self._swap_nodes(__magic_name__ ,			__magic_name__  )
                return self._bubble_down(__magic_name__  )
      if child_left_position < self.elements:
           _a ,				_a      =							self.heap[child_left_position]
           if child_left_weight < weight:
                self._swap_nodes(__magic_name__ ,			__magic_name__  )
                return self._bubble_down(__magic_name__  )
      else:
           return None
      if child_right_position < self.elements:
           _a ,				_a      =							self.heap[child_right_position]
           if child_right_weight < weight:
                self._swap_nodes(__magic_name__ ,			__magic_name__  )
                return self._bubble_down(__magic_name__  )
      return None




 def 							__UpperCAmelCase						(	self ,			__magic_name__ ,			__magic_name__  )   ->							None:
      # Swap the nodes at the given positions
      _a      =							self.heap[nodea_pos][0]
      _a      =							self.heap[nodea_pos][0]
      _a ,				_a      =							(
          self.heap[nodea_pos],
          self.heap[nodea_pos],
      )
      _a      =							nodea_pos
      _a      =							nodea_pos





class 	a       (       Generic[T]    ):
 def __init__(	self  )   ->							None:
      _a      =							{}
      _a      =							0
 def __repr__(	self  )   ->							str:
      return str(self.connections  )
 def __len__(	self  )   ->							int:
      return self.nodes
 def 							__UpperCAmelCase						(	self ,			__magic_name__  )   ->							None:
      # Add a node in the graph if it is not in the graph
      if node not in self.connections:
           _a      =							{}
           self.nodes += 1




 def 							__UpperCAmelCase						(	self ,			__magic_name__ ,			__magic_name__ ,			__magic_name__  )   ->							None:
      # Add an edge between 2 nodes in the graph
      self.add_node(__magic_name__  )
      self.add_node(__magic_name__  )
      _a      =							weight
      _a      =							weight






def 		_A			(lowerCAmelCase__					:GraphUndirectedWeighted[T]			,  )				->					tuple[dict[T, int], dict[T, T | None]]:




     '''simple docstring'''



     _a      =							{node: maxsize for node in graph.connections}
     _a      =							{node: None for node in graph.connections}

     _a      =							MinPriorityQueue()
     for node, weight in dist.items():
          priority_queue.push(lowerCAmelCase__			,  lowerCAmelCase__						)

     if priority_queue.is_empty():
          return dist, parent

     # initialization
     _a      =							priority_queue.extract_min()
     _a      =							0
     for neighbour in graph.connections[node]:
          if dist[neighbour] > dist[node] + graph.connections[node][neighbour]:
               _a      =							dist[node] + graph.connections[node][neighbour]
               priority_queue.update_key(lowerCAmelCase__			,  dist[neighbour]						)
               _a      =							node

    # running prim's algorithm
     while not priority_queue.is_empty():
          _a      =							priority_queue.extract_min()
          for neighbour in graph.connections[node]:
               if dist[neighbour] > dist[node] + graph.connections[node][neighbour]:
                    _a      =							dist[node] + graph.connections[node][neighbour]
                    priority_queue.update_key(lowerCAmelCase__			,  dist[neighbour]						)
                    _a      =							node
     return dist, parent

