



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       A_   (						__a  :					Tuple		):

				"""simple docstring"""
				a__					   =		"""https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"""

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

def       A_   (						):

				"""simple docstring"""
				a__					   =		VivitConfig()

				a__					   =		400
				a__					   =		"""huggingface/label-files"""
				a__					   =		"""kinetics400-id2label.json"""

				a__					   =		json.load(open(hf_hub_download(__a			,		__a			,		repo_type="""dataset"""		)			,		"""r"""		)		)
				a__					   =		{int(__a		): v for k, v in idalabel.items()}
				a__					   =		idalabel
				a__					   =		{v: k for k, v in idalabel.items()}
				return config

def       A_   (						):

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

def       A_   (						__a  :					np.ndarray		):

				"""simple docstring"""
				if np.ndim(__a		) == 2:
								return transform_attention_bias(__a		)

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

				else:
								raise Exception(F'''Invalid number of dimesions: {np.ndim(__a		)}'''		)

def       A_   (						__a  :					np.ndarray		):

				"""simple docstring"""
				return current.flatten()

def       A_   (						__a  :					np.ndarray		):

				"""simple docstring"""
				return np.reshape(__a			,		(current.shape[0], current.shape[1] * current.shape[2])		).T

def       A_   (						__a  :					np.ndarray		):

				"""simple docstring"""
				return np.reshape(__a			,		(current.shape[0] * current.shape[1], current.shape[2])		).T

def       A_   (						__a  :					List[Any]			,		__a  :					int		):

				"""simple docstring"""
				a__					   =		state_dict["""optimizer"""]["""target"""]["""Transformer"""][F'''encoderblock_{i}''']

				a__					   =		F'''encoder.layer.{i}.'''
				a__					   =		{
				    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       A_   (						__a  :					str		):

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

def       A_   (						__a  :					Dict			,		__a  :					Union[str, Any]=False		):

				"""simple docstring"""
				a__					   =		get_n_layers(__a		)

				a__					   =		OrderedDict()

				a__					   =		state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""encoder_norm"""]["""bias"""]
				a__					   =		state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""encoder_norm"""]["""scale"""]

				a__					   =		np.transpose(
				    state_dict["""optimizer"""]["""target"""]["""embedding"""]["""kernel"""]			,		(4, 3, 0, 1, 2)		)
				a__					   =		state_dict["""optimizer"""]["""target"""]["""embedding"""]["""bias"""]

				a__					   =		state_dict["""optimizer"""]["""target"""]["""cls"""]
				a__					   =		state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""posembed_input"""][
				    """pos_embedding"""
				]

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

				if classification_head:
								a__					   =		{"""vivit.""" + k: v for k, v in new_state.items()}
								a__					   =		np.transpose(state_dict["""optimizer"""]["""target"""]["""output_projection"""]["""kernel"""]		)
								a__					   =		np.transpose(state_dict["""optimizer"""]["""target"""]["""output_projection"""]["""bias"""]		)

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

def       A_   (						):

				"""simple docstring"""
				a__					   =		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       A_   (						__a  :					str		):

				"""simple docstring"""
				a__					   =		"""checkpoint"""

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

				a__					   =		restore_checkpoint(__a			,		__a		)
				a__					   =		transform_state(__a			,		classification_head=__a		)

				a__					   =		get_vivit_config()

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

				a__					   =		VivitForVideoClassification(__a		)
				model.load_state_dict(__a		)
				model.eval()

				a__					   =		get_processor()

				a__					   =		prepare_video()
				a__					   =		extractor(__a			,		return_tensors="""pt"""		)

				a__					   =		model(**__a		)

				a__					   =		torch.Size([1, 400]		)
				a__					   =		torch.tensor([-1.0_5_4_3, 2.0_7_6_4, -0.2_1_0_4, 0.4_4_3_9, -0.9_6_5_8]		)

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

				model.save_pretrained(__a		)
				extractor.save_pretrained(__a		)


if __name__ == "__main__":
	UpperCAmelCase							=       argparse.ArgumentParser()

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

	UpperCAmelCase							=       parser.parse_args()
	convert(args.output_model_name)




import sys
from collections.abc import Mapping

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter



class     lowercase_       (    TensorFormatter[Mapping, np.ndarray, Mapping]	):



      def __init__(     self :							Union[str, Any]     ,   A__ :							Dict=None     ,   **A__ :							str						)       ->	Any:
        super().__init__(features=A__						)
        _snake_case		=				np_array_kwargs



      def 		UpperCamelCase_						(     self :							str     ,   A__ :							Dict						)       ->	Dict:
        if isinstance(A__     ,   A__						):
          if column and all(
              isinstance(A__     ,   np.ndarray						) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column						):
            return np.stack(A__						)
          else:
            # don't use np.array(column, dtype=object)
            # since it fails in certain cases
            # see https://stackoverflow.com/q/51005699
            _snake_case		=				np.empty(len(A__						)     ,   dtype=A__						)
            _snake_case		=				column
            return out
        return column



      def 		UpperCamelCase_						(     self :							Tuple     ,   A__ :							Optional[int]						)       ->	Optional[Any]:
        if isinstance(A__     ,   (str, bytes, type(A__						))						):
          return value
        elif isinstance(A__     ,   (np.character, np.ndarray)						) and np.issubdtype(value.dtype     ,   np.character						):
          return value
        elif isinstance(A__     ,   np.number						):
          return value

        _snake_case		=				{}

        if isinstance(A__     ,   np.ndarray						) and np.issubdtype(value.dtype     ,   np.integer						):
          _snake_case		=				{'''dtype''': np.intaa}
        elif isinstance(A__     ,   np.ndarray						) and np.issubdtype(value.dtype     ,   np.floating						):
          _snake_case		=				{'''dtype''': np.floataa}
        elif config.PIL_AVAILABLE and "PIL" in sys.modules:
          import PIL.Image

          if isinstance(A__     ,   PIL.Image.Image						):
            return np.asarray(A__     ,   **self.np_array_kwargs						)

        return np.asarray(A__     ,   **{**default_dtype, **self.np_array_kwargs}						)



      def 		UpperCamelCase_						(     self :							Union[str, Any]     ,   A__ :							Union[str, Any]						)       ->	int:
        # support for torch, tf, jax etc.
        if config.TORCH_AVAILABLE and "torch" in sys.modules:
          import torch

          if isinstance(A__     ,   torch.Tensor						):
            return self._tensorize(data_struct.detach().cpu().numpy()[()]						)
        if hasattr(A__     ,   '''__array__'''						) and not isinstance(A__     ,   (np.ndarray, np.character, np.number)						):
          _snake_case		=				data_struct.__array__()
        # support for nested types like struct of list of struct
        if isinstance(A__     ,   np.ndarray						):
          if data_struct.dtype == object:
            return self._consolidate([self.recursive_tensorize(A__						) for substruct in data_struct]						)
        if isinstance(A__     ,   (list, tuple)						):
          return self._consolidate([self.recursive_tensorize(A__						) for substruct in data_struct]						)
        return self._tensorize(A__						)



      def 		UpperCamelCase_						(     self :							List[Any]     ,   A__ :							dict						)       ->	Tuple:
        return map_nested(self._recursive_tensorize     ,   A__     ,   map_list=A__						)



      def 		UpperCamelCase_						(     self :							str     ,   A__ :							pa.Table						)       ->	Mapping:
        _snake_case		=				self.numpy_arrow_extractor().extract_row(A__						)
        _snake_case		=				self.python_features_decoder.decode_row(A__						)
        return self.recursive_tensorize(A__						)



      def 		UpperCamelCase_						(     self :							Optional[Any]     ,   A__ :							pa.Table						)       ->	np.ndarray:
        _snake_case		=				self.numpy_arrow_extractor().extract_column(A__						)
        _snake_case		=				self.python_features_decoder.decode_column(A__     ,   pa_table.column_names[0]						)
        _snake_case		=				self.recursive_tensorize(A__						)
        _snake_case		=				self._consolidate(A__						)
        return column







      def 		UpperCamelCase_						(     self :							Dict     ,   A__ :							pa.Table						)       ->	Mapping:
        _snake_case		=				self.numpy_arrow_extractor().extract_batch(A__						)
        _snake_case		=				self.python_features_decoder.decode_batch(A__						)
        _snake_case		=				self.recursive_tensorize(A__						)
        for column_name in batch:
          _snake_case		=				self._consolidate(batch[column_name]						)
        return batch

