code: 
  
  
  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     UpperCAmelCase__      (			UpperCAmelCase_					:       Dict							)  ->	int:
       __lowerCamelCase     :			Any              =	'https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'
  
       with open(UpperCAmelCase_   ,					'wb'							) as f:
            with requests.get(UpperCAmelCase_   ,					stream=UpperCAmelCase_							) as req:
                 for chunk in req.iter_content(chunk_size=20_48							):
                      f.write(UpperCAmelCase_							)
  
  
  
  def     UpperCAmelCase__      (			)  ->	VivitConfig:
       __lowerCamelCase     :			Union[str, Any]              =	VivitConfig()
  
       __lowerCamelCase     :			Optional[int]              =	4_00
       __lowerCamelCase     :			str              =	'huggingface/label-files'
       __lowerCamelCase     :			Optional[Any]              =	'kinetics400-id2label.json'
  
       __lowerCamelCase     :			Tuple              =	json.load(open(hf_hub_download(UpperCAmelCase_   ,					UpperCAmelCase_   ,					repo_type='dataset'							)   ,					'r'							)							)
       __lowerCamelCase     :			Optional[Any]              =	{int(UpperCAmelCase_							): v for k, v in idalabel.items()}
       __lowerCamelCase     :			int              =	idalabel
       __lowerCamelCase     :			Any              =	{v: k for k, v in idalabel.items()}
       return config
  
  
  
  def     UpperCAmelCase__      (			)  ->	Optional[int]:
       __lowerCamelCase     :			Optional[int]              =	hf_hub_download(
           repo_id='hf-internal-testing/spaghetti-video'   ,					filename='eating_spaghetti_32_frames.npy'   ,					repo_type='dataset'							)
       __lowerCamelCase     :			Optional[int]              =	np.load(UpperCAmelCase_							)
       return list(UpperCAmelCase_							)
  
  
  
  def     UpperCAmelCase__      (			UpperCAmelCase_					:       np.ndarray							)  ->	Dict:
       if np.ndim(UpperCAmelCase_							) == 2:
            return transform_attention_bias(UpperCAmelCase_							)
  
       elif np.ndim(UpperCAmelCase_							) == 3:
            return transform_attention_kernel(UpperCAmelCase_							)
  
       else:
            raise Exception(F'Invalid number of dimesions: {np.ndim(UpperCAmelCase_							)}'							)
  
  
  
  def     UpperCAmelCase__      (			UpperCAmelCase_					:       np.ndarray							)  ->	List[Any]:
       return current.flatten()
  
  
  
  def     UpperCAmelCase__      (			UpperCAmelCase_					:       np.ndarray							)  ->	str:
       return np.reshape(UpperCAmelCase_   ,					(current.shape[0], current.shape[1] * current.shape[2])							).T
  
  
  
  def     UpperCAmelCase__      (			UpperCAmelCase_					:       np.ndarray							)  ->	List[Any]:
       return np.reshape(UpperCAmelCase_   ,					(current.shape[0] * current.shape[1], current.shape[2])							).T
  
  
  
  def     UpperCAmelCase__      (			UpperCAmelCase_					:       int   ,					UpperCAmelCase_					:       Tuple							)  ->	Any:
       __lowerCamelCase     :			Union[str, Any]              =	state_dict['optimizer']['target']['Transformer'][F'encoderblock_{i}']
  
       __lowerCamelCase     :			Dict              =	F'encoder.layer.{i}.'
       __lowerCamelCase     :			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     UpperCAmelCase__      (			UpperCAmelCase_					:       int							)  ->	Optional[int]:
       return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]							)
  
  
  
  def     UpperCAmelCase__      (			UpperCAmelCase_					:       List[str]   ,					UpperCAmelCase_					:       int=False							)  ->	Optional[int]:
       __lowerCamelCase     :			Tuple              =	get_n_layers(UpperCAmelCase_							)
  
       __lowerCamelCase     :			int              =	OrderedDict()
  
       __lowerCamelCase     :			Union[str, Any]              =	state_dict['optimizer']['target']['Transformer']['encoder_norm']['bias']
       __lowerCamelCase     :			Optional[Any]              =	state_dict['optimizer']['target']['Transformer']['encoder_norm']['scale']
  
       __lowerCamelCase     :			Union[str, Any]              =	np.transpose(
           state_dict['optimizer']['target']['embedding']['kernel']   ,					(4, 3, 0, 1, 2)							)
       __lowerCamelCase     :			str              =	state_dict['optimizer']['target']['embedding']['bias']
  
       __lowerCamelCase     :			List[str]              =	state_dict['optimizer']['target']['cls']
       __lowerCamelCase     :			int              =	state_dict['optimizer']['target']['Transformer']['posembed_input'][
           'pos_embedding'
       ]
  
       for i in range(UpperCAmelCase_							):
            new_state.update(transform_state_encoder_block(UpperCAmelCase_   ,					UpperCAmelCase_							)							)
  
       if classification_head:
            __lowerCamelCase     :			List[Any]              =	{'vivit.' + k: v for k, v in new_state.items()}
            __lowerCamelCase     :			List[Any]              =	np.transpose(state_dict['optimizer']['target']['output_projection']['kernel']							)
            __lowerCamelCase     :			Optional[int]              =	np.transpose(state_dict['optimizer']['target']['output_projection']['bias']							)
  
       return {k: torch.tensor(UpperCAmelCase_							) for k, v in new_state.items()}
  
  
  
  def     UpperCAmelCase__      (			)  ->	VivitImageProcessor:
       __lowerCamelCase     :			List[Any]              =	VivitImageProcessor()
  
       assert extractor.do_resize is True
       assert extractor.size == {"shortest_edge": 2_56}
       assert extractor.do_center_crop is True
       assert extractor.crop_size == {"width": 2_24, "height": 2_24}
       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 / 2_55
  
       # zero-centering = True in original implementation
       assert extractor.do_zero_centering is True
  
       return extractor
  
  
  
  def     UpperCAmelCase__      (			UpperCAmelCase_					:       str							)  ->	Tuple:
       __lowerCamelCase     :			int              =	'checkpoint'
  
       if not os.path.exists(UpperCAmelCase_							):
            download_checkpoint(UpperCAmelCase_							)
  
       __lowerCamelCase     :			int              =	restore_checkpoint(UpperCAmelCase_   ,					UpperCAmelCase_							)
       __lowerCamelCase     :			Union[str, Any]              =	transform_state(UpperCAmelCase_   ,					classification_head=UpperCAmelCase_							)
  
       __lowerCamelCase     :			Union[str, Any]              =	get_vivit_config()
  
       assert config.image_size == 2_24
       assert config.num_frames == 32
  
       __lowerCamelCase     :			int              =	VivitForVideoClassification(UpperCAmelCase_							)
       model.load_state_dict(UpperCAmelCase_							)
       model.eval()
  
       __lowerCamelCase     :			Optional[Any]              =	get_processor()
  
       __lowerCamelCase     :			Tuple              =	prepare_video()
       __lowerCamelCase     :			Any              =	extractor(UpperCAmelCase_   ,					return_tensors='pt'							)
  
       __lowerCamelCase     :			Tuple              =	model(**UpperCAmelCase_							)
  
       __lowerCamelCase     :			Union[str, Any]              =	torch.Size([1, 4_00]							)
       __lowerCamelCase     :			int              =	torch.tensor([-1.0_543, 2.0_764, -0.2_104, 0.4_439, -0.9_658]							)
  
       assert outputs.logits.shape == expected_shape
       assert torch.allclose(outputs.logits[0, :5]   ,					UpperCAmelCase_   ,					atol=1e-4							), outputs.logits[0, :5]
  
       model.save_pretrained(UpperCAmelCase_							)
       extractor.save_pretrained(UpperCAmelCase_							)
  
  
  if __name__ == "__main__":
        A__	:      Union[str, Any]				=						argparse.ArgumentParser()
  
        parser.add_argument("""--output_model_name""", """-o""", type=str, help="""Output path for the converted HuggingFace model""")
  
        A__	:      List[str]				=						parser.parse_args()
        convert(args.output_model_name)
  
  
  
  
  
  
  
style_context: 
  
  
  '''simple docstring'''
  
  
  
  from typing import TYPE_CHECKING
  
  from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available
  
  
  A__	:      Optional[Any]				=						{
      """configuration_xlm_prophetnet""": ["""XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP""", """XLMProphetNetConfig"""],
  }
  
  try:
        if not is_sentencepiece_available():
              raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
        pass
  else:
        A__	:      List[Any]				=						["""XLMProphetNetTokenizer"""]
  
  try:
        if not is_torch_available():
              raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
        pass
  else:
        A__	:      Optional[Any]				=						[
            """XLM_PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST""",
            """XLMProphetNetDecoder""",
            """XLMProphetNetEncoder""",
            """XLMProphetNetForCausalLM""",
            """XLMProphetNetForConditionalGeneration""",
            """XLMProphetNetModel""",
            """XLMProphetNetPreTrainedModel""",
        ]
  
  
  if TYPE_CHECKING:
        from .configuration_xlm_prophetnet import XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMProphetNetConfig
  
        try:
              if not is_sentencepiece_available():
                    raise OptionalDependencyNotAvailable()
        except OptionalDependencyNotAvailable:
              pass
        else:
              from .tokenization_xlm_prophetnet import XLMProphetNetTokenizer
  
        try:
              if not is_torch_available():
                    raise OptionalDependencyNotAvailable()
        except OptionalDependencyNotAvailable:
              pass
        else:
              from .modeling_xlm_prophetnet import (
                  XLM_PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST,
                  XLMProphetNetDecoder,
                  XLMProphetNetEncoder,
                  XLMProphetNetForCausalLM,
                  XLMProphetNetForConditionalGeneration,
                  XLMProphetNetModel,
                  XLMProphetNetPreTrainedModel,
              )
  
  else:
        import sys
  
        A__	:      List[str]				=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
  
  
  
  
  
  
  