


"""simple docstring"""
import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling






def 						lowercase				(	_snake_case		: Optional[Any]	)	->Union[str, Any]:
	"""simple docstring"""





	__snake_case      :		str         =							'''https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'''

	with open(_snake_case	,  '''wb'''	) as f:
		with requests.get(_snake_case	,  stream=_snake_case	) as req:
			for chunk in req.iter_content(chunk_size=2_048	):
				f.write(_snake_case	)






def 						lowercase				(	)	->VivitConfig:
	"""simple docstring"""





	__snake_case      :		int         =							VivitConfig()

	__snake_case      :		Tuple         =							400
	__snake_case      :		Tuple         =							'''huggingface/label-files'''
	__snake_case      :		List[str]         =							'''kinetics400-id2label.json'''

	__snake_case      :		Dict         =							json.load(open(hf_hub_download(_snake_case	,  _snake_case	,  repo_type='''dataset'''	)	,  '''r'''	)	)
	__snake_case      :		List[str]         =							{int(_snake_case	): v for k, v in idalabel.items()}
	__snake_case      :		Tuple         =							idalabel
	__snake_case      :		Optional[int]         =							{v: k for k, v in idalabel.items()}
	return config






def 						lowercase				(	)	->Optional[int]:
	"""simple docstring"""





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






def 						lowercase				(	_snake_case		: np.ndarray	)	->str:
	"""simple docstring"""





	if np.ndim(_snake_case	) == 2:
		return transform_attention_bias(_snake_case	)

	elif np.ndim(_snake_case	) == 3:
		return transform_attention_kernel(_snake_case	)

	else:
		raise Exception(f"""Invalid number of dimesions: {np.ndim(_snake_case	)}"""	)






def 						lowercase				(	_snake_case		: np.ndarray	)	->Tuple:
	"""simple docstring"""





	return current.flatten()






def 						lowercase				(	_snake_case		: np.ndarray	)	->List[str]:
	"""simple docstring"""





	return np.reshape(_snake_case	,  (current.shape[0], current.shape[1] * current.shape[2])	).T






def 						lowercase				(	_snake_case		: np.ndarray	)	->Optional[Any]:
	"""simple docstring"""





	return np.reshape(_snake_case	,  (current.shape[0] * current.shape[1], current.shape[2])	).T






def 						lowercase				(	_snake_case		: Dict	,  _snake_case		: Tuple	)	->Tuple:
	"""simple docstring"""





	__snake_case      :		int         =							state_dict['''optimizer''']['''target''']['''Transformer'''][f"""encoderblock_{i}"""]

	__snake_case      :		Optional[int]         =							f"""encoder.layer.{i}."""
	__snake_case      :		Dict         =							{
	    prefix + '''intermediate.dense.bias''': state['''MlpBlock_0''']['''Dense_0''']['''bias'''],
	    prefix + '''intermediate.dense.weight''': np.transpose(state['''MlpBlock_0''']['''Dense_0''']['''kernel''']	),
	    prefix + '''output.dense.bias''': state['''MlpBlock_0''']['''Dense_1''']['''bias'''],
	    prefix + '''output.dense.weight''': np.transpose(state['''MlpBlock_0''']['''Dense_1''']['''kernel''']	),
	    prefix + '''layernorm_before.bias''': state['''LayerNorm_0''']['''bias'''],
	    prefix + '''layernorm_before.weight''': state['''LayerNorm_0''']['''scale'''],
	    prefix + '''layernorm_after.bias''': state['''LayerNorm_1''']['''bias'''],
	    prefix + '''layernorm_after.weight''': state['''LayerNorm_1''']['''scale'''],
	    prefix
	    + '''attention.attention.query.bias''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''query''']['''bias''']	),
	    prefix
	    + '''attention.attention.query.weight''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''query''']['''kernel''']	),
	    prefix
	    + '''attention.attention.key.bias''': transform_attention(state['''MultiHeadDotProductAttention_0''']['''key''']['''bias''']	),
	    prefix
	    + '''attention.attention.key.weight''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''key''']['''kernel''']	),
	    prefix
	    + '''attention.attention.value.bias''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''value''']['''bias''']	),
	    prefix
	    + '''attention.attention.value.weight''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''value''']['''kernel''']	),
	    prefix + '''attention.output.dense.bias''': state['''MultiHeadDotProductAttention_0''']['''out''']['''bias'''],
	    prefix
	    + '''attention.output.dense.weight''': transform_attention_output_weight(
	        state['''MultiHeadDotProductAttention_0''']['''out''']['''kernel''']	),
	}

	return new_state






def 						lowercase				(	_snake_case		: int	)	->List[str]:
	"""simple docstring"""





	return sum([1 if '''encoderblock_''' in k else 0 for k in state_dict['''optimizer''']['''target''']['''Transformer'''].keys()]	)






def 						lowercase				(	_snake_case		: Optional[Any]	,  _snake_case		: Optional[int]=False	)	->Union[str, Any]:
	"""simple docstring"""





	__snake_case      :		Tuple         =							get_n_layers(_snake_case	)

	__snake_case      :		str         =							OrderedDict()

	__snake_case      :		List[Any]         =							state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''bias''']
	__snake_case      :		List[str]         =							state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''scale''']

	__snake_case      :		Optional[Any]         =							np.transpose(
	    state_dict['''optimizer''']['''target''']['''embedding''']['''kernel''']	,  (4, 3, 0, 1, 2)	)
	__snake_case      :		Optional[Any]         =							state_dict['''optimizer''']['''target''']['''embedding''']['''bias''']

	__snake_case      :		Union[str, Any]         =							state_dict['''optimizer''']['''target''']['''cls''']
	__snake_case      :		List[str]         =							state_dict['''optimizer''']['''target''']['''Transformer''']['''posembed_input'''][
	    '''pos_embedding'''
	]

	for i in range(_snake_case	):
		new_state.update(transform_state_encoder_block(_snake_case	,  _snake_case	)	)

	if classification_head:
		__snake_case      :		List[Any]         =							{'''vivit.''' + k: v for k, v in new_state.items()}
		__snake_case      :		str         =							np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''kernel''']	)
		__snake_case      :		int         =							np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''bias''']	)

	return {k: torch.tensor(_snake_case	) for k, v in new_state.items()}






def 						lowercase				(	)	->VivitImageProcessor:
	"""simple docstring"""





	__snake_case      :		List[str]         =							VivitImageProcessor()

	assert extractor.do_resize is True
	assert extractor.size == {"shortest_edge": 256}
	assert extractor.do_center_crop is True
	assert extractor.crop_size == {"width": 224, "height": 224}
	assert extractor.resample == PILImageResampling.BILINEAR

	# here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
	# one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
	# which effectively means no normalization (and ViViT does not overwrite those when calling this func)
	assert extractor.do_normalize is False
	assert extractor.do_rescale is True
	assert extractor.rescale_factor == 1 / 255

	# zero-centering = True in original implementation
	assert extractor.do_zero_centering is True

	return extractor






def 						lowercase				(	_snake_case		: str	)	->Optional[Any]:
	"""simple docstring"""





	__snake_case      :		Tuple         =							'''checkpoint'''

	if not os.path.exists(_snake_case	):
		download_checkpoint(_snake_case	)

	__snake_case      :		List[Any]         =							restore_checkpoint(_snake_case	,  _snake_case	)
	__snake_case      :		Union[str, Any]         =							transform_state(_snake_case	,  classification_head=_snake_case	)

	__snake_case      :		Optional[Any]         =							get_vivit_config()

	assert config.image_size == 224
	assert config.num_frames == 32

	__snake_case      :		List[str]         =							VivitForVideoClassification(_snake_case	)
	model.load_state_dict(_snake_case	)
	model.eval()

	__snake_case      :		Tuple         =							get_processor()

	__snake_case      :		Any         =							prepare_video()
	__snake_case      :		List[str]         =							extractor(_snake_case	,  return_tensors='''pt'''	)

	__snake_case      :		Union[str, Any]         =							model(**_snake_case	)

	__snake_case      :		List[str]         =							torch.Size([1, 400]	)
	__snake_case      :		Optional[int]         =							torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]	)

	assert outputs.logits.shape == expected_shape
	assert torch.allclose(outputs.logits[0, :5]	,  _snake_case	,  atol=1e-4	), outputs.logits[0, :5]

	model.save_pretrained(_snake_case	)
	extractor.save_pretrained(_snake_case	)


if __name__ == "__main__":
			SCREAMING_SNAKE_CASE					:       Optional[Any]							   =							argparse.ArgumentParser()

			parser.add_argument("""--output_model_name""", """-o""", type=str, help="""Output path for the converted HuggingFace model""")

			SCREAMING_SNAKE_CASE					:       List[str]							   =							parser.parse_args()
			convert(args.output_model_name)









import argparse
import os
import re


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_dummies.py
__UpperCamelCase							:							Optional[Any]    							=					"src/transformers"

# Matches is_xxx_available()
__UpperCamelCase							:							str    							=					re.compile(r"is\_([a-z_]*)_available()")
# Matches from xxx import bla
__UpperCamelCase							:							Union[str, Any]    							=					re.compile(r"\s+from\s+\S*\s+import\s+([^\(\s].*)\n")
__UpperCamelCase							:							Dict    							=					re.compile(r"^\s+if\s+not\s+\(?is\_[a-z_]*\_available\(\)")


__UpperCamelCase							:							Any    							=					"\n{0} = None\n"

__UpperCamelCase							:							Tuple    							=					"\nclass {0}(metaclass=DummyObject):\n    _backends = {1}\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, {1})\n"


__UpperCamelCase							:							Optional[Any]    							=					"\ndef {0}(*args, **kwargs):\n    requires_backends({0}, {1})\n"

def       _a	(					SCREAMING_SNAKE_CASE						:					str  ):







 """simple docstring"""



 if _re_test_backend.search(SCREAMING_SNAKE_CASE  ) is None:
  return None
 UpperCamelCase__       :     List[str] 							= [b[0] for b in _re_backend.findall(SCREAMING_SNAKE_CASE  )]
 backends.sort()
 return "_and_".join(SCREAMING_SNAKE_CASE  )

def       _a	(					):







 """simple docstring"""



 with open(os.path.join(SCREAMING_SNAKE_CASE     ,			'''__init__.py'''  )     ,			'''r'''     ,			encoding='''utf-8'''     ,			newline='''\n'''  ) as f:
  UpperCamelCase__       :     Union[str, Any] 							= f.readlines()

 # Get to the point we do the actual imports for type checking
 UpperCamelCase__       :     int 							= 0
 while not lines[line_index].startswith('''if TYPE_CHECKING'''  ):
  line_index += 1

 UpperCamelCase__       :     Optional[Any] 							= {}
 # Go through the end of the file
 while line_index < len(SCREAMING_SNAKE_CASE  ):
  # If the line is an if is_backend_available, we grab all objects associated.
  UpperCamelCase__       :     Any 							= find_backend(lines[line_index]  )
  if backend is not None:
   while not lines[line_index].startswith('''    else:'''  ):
    line_index += 1
   line_index += 1

   UpperCamelCase__       :     List[str] 							= []
   # Until we unindent, add backend objects to the list
   while len(lines[line_index]  ) <= 1 or lines[line_index].startswith(''' ''' * 8  ):
    UpperCamelCase__       :     str 							= lines[line_index]
    UpperCamelCase__       :     Union[str, Any] 							= _re_single_line_import.search(SCREAMING_SNAKE_CASE  )
    if single_line_import_search is not None:
     objects.extend(single_line_import_search.groups()[0].split(''', '''  )  )
    elif line.startswith(''' ''' * 12  ):
     objects.append(line[12:-2]  )
    line_index += 1

   UpperCamelCase__       :     Optional[int] 							= objects
  else:
   line_index += 1

 return backend_specific_objects

def       _a	(					SCREAMING_SNAKE_CASE						:					str     ,			SCREAMING_SNAKE_CASE						:					List[str]  ):







 """simple docstring"""



 if name.isupper():
  return DUMMY_CONSTANT.format(SCREAMING_SNAKE_CASE  )
 elif name.islower():
  return DUMMY_FUNCTION.format(SCREAMING_SNAKE_CASE     ,			SCREAMING_SNAKE_CASE  )
 else:
  return DUMMY_CLASS.format(SCREAMING_SNAKE_CASE     ,			SCREAMING_SNAKE_CASE  )

def       _a	(					SCREAMING_SNAKE_CASE						:					Dict=None  ):







 """simple docstring"""



 if backend_specific_objects is None:
  UpperCamelCase__       :     Any 							= read_init()
 # For special correspondence backend to module name as used in the function requires_modulename
 UpperCamelCase__       :     Optional[Any] 							= {}

 for backend, objects in backend_specific_objects.items():
  UpperCamelCase__       :     int 							= '''[''' + ''', '''.join(F"\"{b}\"" for b in backend.split('''_and_'''  )  ) + ''']'''
  UpperCamelCase__       :     Any 							= '''# This file is autogenerated by the command `make fix-copies`, do not edit.\n'''
  dummy_file += "from ..utils import DummyObject, requires_backends\n\n"
  dummy_file += "\n".join([create_dummy_object(SCREAMING_SNAKE_CASE     ,			SCREAMING_SNAKE_CASE  ) for o in objects]  )
  UpperCamelCase__       :     List[str] 							= dummy_file

 return dummy_files

def       _a	(					SCREAMING_SNAKE_CASE						:					Union[str, Any]=False  ):







 """simple docstring"""



 UpperCamelCase__       :     Dict 							= create_dummy_files()
 # For special correspondence backend to shortcut as used in utils/dummy_xxx_objects.py
 UpperCamelCase__       :     Tuple 							= {'''torch''': '''pt'''}

 # Locate actual dummy modules and read their content.
 UpperCamelCase__       :     Optional[Any] 							= os.path.join(SCREAMING_SNAKE_CASE     ,			'''utils'''  )
 UpperCamelCase__       :     int 							= {
     backend: os.path.join(SCREAMING_SNAKE_CASE     ,			F"dummy_{short_names.get(SCREAMING_SNAKE_CASE     ,			SCREAMING_SNAKE_CASE  )}_objects.py"  )
     for backend in dummy_files.keys()
 }

 UpperCamelCase__       :     Optional[Any] 							= {}
 for backend, file_path in dummy_file_paths.items():
  if os.path.isfile(SCREAMING_SNAKE_CASE  ):
   with open(SCREAMING_SNAKE_CASE     ,			'''r'''     ,			encoding='''utf-8'''     ,			newline='''\n'''  ) as f:
    UpperCamelCase__       :     List[Any] 							= f.read()
  else:
   UpperCamelCase__       :     List[Any] 							= ''''''

 for backend in dummy_files.keys():
  if dummy_files[backend] != actual_dummies[backend]:
   if overwrite:
    print(
        F"Updating transformers.utils.dummy_{short_names.get(SCREAMING_SNAKE_CASE     ,			SCREAMING_SNAKE_CASE  )}_objects.py as the main "
        '''__init__ has new objects.'''  )
    with open(dummy_file_paths[backend]     ,			'''w'''     ,			encoding='''utf-8'''     ,			newline='''\n'''  ) as f:
     f.write(dummy_files[backend]  )
   else:
    raise ValueError(
        '''The main __init__ has objects that are not present in '''
        F"transformers.utils.dummy_{short_names.get(SCREAMING_SNAKE_CASE     ,			SCREAMING_SNAKE_CASE  )}_objects.py. Run `make fix-copies` "
        '''to fix this.'''  )


if __name__ == "__main__":
       __UpperCamelCase							:							Optional[int]    							=					argparse.ArgumentParser()
       parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
       __UpperCamelCase							:							str    							=					parser.parse_args()

       check_dummies(args.fix_and_overwrite)


