
"""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 				a__			(					lowerCAmelCase						:							List[Any]		):






     '''simple docstring'''
     UpperCAmelCase__				:			List[str]							=    "https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"

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



def 				a__			(					):






     '''simple docstring'''
     UpperCAmelCase__				:			int							=    VivitConfig()

     UpperCAmelCase__				:			Dict							=    400
     UpperCAmelCase__				:			Tuple							=    "huggingface/label-files"
     UpperCAmelCase__				:			Tuple							=    "kinetics400-id2label.json"

     UpperCAmelCase__				:			Tuple							=    json.load(open(hf_hub_download(lowerCAmelCase							,     lowerCAmelCase							,     repo_type="dataset"		)							,     "r"		)		)
     UpperCAmelCase__				:			List[Any]							=    {int(lowerCAmelCase		): v for k, v in idalabel.items()}
     UpperCAmelCase__				:			Dict							=    idalabel
     UpperCAmelCase__				:			str							=    {v: k for k, v in idalabel.items()}
     return config



def 				a__			(					):






     '''simple docstring'''
     UpperCAmelCase__				:			List[str]							=    hf_hub_download(
         repo_id="hf-internal-testing/spaghetti-video"							,     filename="eating_spaghetti_32_frames.npy"							,     repo_type="dataset"		)
     UpperCAmelCase__				:			Tuple							=    np.load(lowerCAmelCase		)
     return list(lowerCAmelCase		)



def 				a__			(					lowerCAmelCase						:							np.ndarray		):






     '''simple docstring'''
     if np.ndim(lowerCAmelCase		) == 2:
          return transform_attention_bias(lowerCAmelCase		)

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

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



def 				a__			(					lowerCAmelCase						:							np.ndarray		):






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



def 				a__			(					lowerCAmelCase						:							np.ndarray		):






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



def 				a__			(					lowerCAmelCase						:							np.ndarray		):






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



def 				a__			(					lowerCAmelCase						:							Dict							,     lowerCAmelCase						:							List[Any]		):






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

     UpperCAmelCase__				:			int							=    F"encoder.layer.{i}."
     UpperCAmelCase__				:			Union[str, Any]							=    {
         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__			(					lowerCAmelCase						:							int		):






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



def 				a__			(					lowerCAmelCase						:							List[str]							,     lowerCAmelCase						:							List[Any]=False		):






     '''simple docstring'''
     UpperCAmelCase__				:			str							=    get_n_layers(lowerCAmelCase		)

     UpperCAmelCase__				:			Optional[int]							=    OrderedDict()

     UpperCAmelCase__				:			Optional[int]							=    state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["bias"]
     UpperCAmelCase__				:			Optional[int]							=    state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["scale"]

     UpperCAmelCase__				:			List[str]							=    np.transpose(
         state_dict["optimizer"]["target"]["embedding"]["kernel"]							,     (4, 3, 0, 1, 2)		)
     UpperCAmelCase__				:			Any							=    state_dict["optimizer"]["target"]["embedding"]["bias"]

     UpperCAmelCase__				:			List[Any]							=    state_dict["optimizer"]["target"]["cls"]
     UpperCAmelCase__				:			Optional[Any]							=    state_dict["optimizer"]["target"]["Transformer"]["posembed_input"][
         "pos_embedding"
     ]

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

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

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



def 				a__			(					):






     '''simple docstring'''
     UpperCAmelCase__				:			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 				a__			(					lowerCAmelCase						:							str		):






     '''simple docstring'''
     UpperCAmelCase__				:			Union[str, Any]							=    "checkpoint"

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

     UpperCAmelCase__				:			Optional[Any]							=    restore_checkpoint(lowerCAmelCase							,     lowerCAmelCase		)
     UpperCAmelCase__				:			List[Any]							=    transform_state(lowerCAmelCase							,     classification_head=lowerCAmelCase		)

     UpperCAmelCase__				:			Optional[int]							=    get_vivit_config()

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

     UpperCAmelCase__				:			Optional[int]							=    VivitForVideoClassification(lowerCAmelCase		)
     model.load_state_dict(lowerCAmelCase		)
     model.eval()

     UpperCAmelCase__				:			Union[str, Any]							=    get_processor()

     UpperCAmelCase__				:			List[Any]							=    prepare_video()
     UpperCAmelCase__				:			str							=    extractor(lowerCAmelCase							,     return_tensors="pt"		)

     UpperCAmelCase__				:			List[Any]							=    model(**lowerCAmelCase		)

     UpperCAmelCase__				:			List[str]							=    torch.Size([1, 400]		)
     UpperCAmelCase__				:			Tuple							=    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]							,     lowerCAmelCase							,     atol=1E-4		), outputs.logits[0, :5]

     model.save_pretrained(lowerCAmelCase		)
     extractor.save_pretrained(lowerCAmelCase		)


if __name__ == "__main__":
    A__   :						Optional[int] 							=     argparse.ArgumentParser()

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

    A__   :						Union[str, Any] 							=     parser.parse_args()
    convert(args.output_model_name)





"""simple docstring"""




def 				a__			(					lowerCAmelCase						:							str							,     lowerCAmelCase						:							int = 13		):






     '''simple docstring'''
     UpperCAmelCase__				:			Any							=    ""
     for c in s:
          if "A" <= c <= "Z":
               out += chr(ord("A"		) + (ord(lowerCAmelCase		) - ord("A"		) + n) % 26		)
          elif "a" <= c <= "z":
               out += chr(ord("a"		) + (ord(lowerCAmelCase		) - ord("a"		) + n) % 26		)
          else:
               out += c
     return out



def 				a__			(					):






     '''simple docstring'''
     UpperCAmelCase__				:			Optional[int]							=    input("Enter message: "		)

     UpperCAmelCase__				:			Optional[Any]							=    dencrypt(lowerCAmelCase							,     13		)
     print("Encryption:"							,     lowerCAmelCase		)

     UpperCAmelCase__				:			List[str]							=    dencrypt(lowerCAmelCase							,     13		)
     print("Decryption: "							,     lowerCAmelCase		)


if __name__ == "__main__":
    import doctest

    doctest.testmod()
    main()



