


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  __SCREAMING_SNAKE_CASE       (				lowercase__  :Any				)       ->  Any:
						__magic_name__				:							List[Any]				=      'https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'

						with open(lowercase__      ,						'wb'				) as f:
												with requests.get(lowercase__      ,						stream=lowercase__				) as req:
																		for chunk in req.iter_content(chunk_size=2_0_4_8				):
																								f.write(lowercase__				)
def  __SCREAMING_SNAKE_CASE       (				)       ->  VivitConfig:
						__magic_name__				:							Dict				=      VivitConfig()

						__magic_name__				:							Any				=      4_0_0
						__magic_name__				:							List[Any]				=      'huggingface/label-files'
						__magic_name__				:							Union[str, Any]				=      'kinetics400-id2label.json'

						__magic_name__				:							int				=      json.load(open(hf_hub_download(lowercase__      ,						lowercase__      ,						repo_type='dataset'				)      ,						'r'				)				)
						__magic_name__				:							Optional[Any]				=      {int(lowercase__				): v for k, v in idalabel.items()}
						__magic_name__				:							List[str]				=      idalabel
						__magic_name__				:							Optional[int]				=      {v: k for k, v in idalabel.items()}
						return config
def  __SCREAMING_SNAKE_CASE       (				)       ->  Tuple:
						__magic_name__				:							int				=      hf_hub_download(
						    repo_id='hf-internal-testing/spaghetti-video'      ,						filename='eating_spaghetti_32_frames.npy'      ,						repo_type='dataset'				)
						__magic_name__				:							List[str]				=      np.load(lowercase__				)
						return list(lowercase__				)
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :np.ndarray				)       ->  Dict:
						if np.ndim(lowercase__				) == 2:
												return transform_attention_bias(lowercase__				)

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

						else:
												raise Exception(f"Invalid number of dimesions: {np.ndim(lowercase__				)}"				)
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :np.ndarray				)       ->  List[Any]:
						return current.flatten()
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :np.ndarray				)       ->  List[str]:
						return np.reshape(lowercase__      ,						(current.shape[0], current.shape[1] * current.shape[2])				).T
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :np.ndarray				)       ->  str:
						return np.reshape(lowercase__      ,						(current.shape[0] * current.shape[1], current.shape[2])				).T
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :Optional[int]      ,						lowercase__  :str				)       ->  Optional[int]:
						__magic_name__				:							Optional[int]				=      state_dict['optimizer']['target']['Transformer'][f"encoderblock_{i}"]

						__magic_name__				:							List[Any]				=      f"encoder.layer.{i}."
						__magic_name__				:							List[str]				=      {
						    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  __SCREAMING_SNAKE_CASE       (				lowercase__  :Tuple				)       ->  Any:
						return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]				)
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :Optional[Any]      ,						lowercase__  :List[str]=False				)       ->  Optional[int]:
						__magic_name__				:							Optional[Any]				=      get_n_layers(lowercase__				)

						__magic_name__				:							List[str]				=      OrderedDict()

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

						__magic_name__				:							Tuple				=      np.transpose(
						    state_dict['optimizer']['target']['embedding']['kernel']      ,						(4, 3, 0, 1, 2)				)
						__magic_name__				:							str				=      state_dict['optimizer']['target']['embedding']['bias']

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

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

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

						return {k: torch.tensor(lowercase__				) for k, v in new_state.items()}
def  __SCREAMING_SNAKE_CASE       (				)       ->  VivitImageProcessor:
						__magic_name__				:							str				=      VivitImageProcessor()

						assert extractor.do_resize is True
						assert extractor.size == {"shortest_edge": 2_5_6}
						assert extractor.do_center_crop is True
						assert extractor.crop_size == {"width": 2_2_4, "height": 2_2_4}
						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_5_5

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

						return extractor
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :str				)       ->  Tuple:
						__magic_name__				:							Dict				=      'checkpoint'

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

						__magic_name__				:							Dict				=      restore_checkpoint(lowercase__      ,						lowercase__				)
						__magic_name__				:							Optional[int]				=      transform_state(lowercase__      ,						classification_head=lowercase__				)

						__magic_name__				:							Any				=      get_vivit_config()

						assert config.image_size == 2_2_4
						assert config.num_frames == 3_2

						__magic_name__				:							List[str]				=      VivitForVideoClassification(lowercase__				)
						model.load_state_dict(lowercase__				)
						model.eval()

						__magic_name__				:							List[Any]				=      get_processor()

						__magic_name__				:							int				=      prepare_video()
						__magic_name__				:							Optional[int]				=      extractor(lowercase__      ,						return_tensors='pt'				)

						__magic_name__				:							Any				=      model(**lowercase__				)

						__magic_name__				:							Optional[Any]				=      torch.Size([1, 4_0_0]				)
						__magic_name__				:							Optional[int]				=      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]      ,						lowercase__      ,						atol=1e-4				), outputs.logits[0, :5]

						model.save_pretrained(lowercase__				)
						extractor.save_pretrained(lowercase__				)


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

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

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






import argparse

from . import (
    BertConfig,
    BertGenerationConfig,
    BertGenerationDecoder,
    BertGenerationEncoder,
    load_tf_weights_in_bert_generation,
    logging,
)


logging.set_verbosity_info()
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :Optional[int]      ,						lowercase__  :List[Any]      ,						lowercase__  :str      ,						lowercase__  :List[str]      ,						lowercase__  :List[Any]				)       ->  List[Any]:
						# Initialise PyTorch model
						__magic_name__				:							Any				=      BertConfig.from_pretrained(
						    'bert-large-cased'      ,						vocab_size=lowercase__      ,						max_position_embeddings=5_1_2      ,						is_decoder=lowercase__      ,						add_cross_attention=lowercase__      ,						)
						__magic_name__				:							Optional[int]				=      bert_config.to_dict()
						del bert_config_dict["type_vocab_size"]
						__magic_name__				:							Optional[int]				=      BertGenerationConfig(**lowercase__				)
						if is_encoder:
												__magic_name__				:							Dict				=      BertGenerationEncoder(lowercase__				)
						else:
												__magic_name__				:							str				=      BertGenerationDecoder(lowercase__				)
						print(f"Building PyTorch model from configuration: {config}"				)

						# Load weights from tf checkpoint
						load_tf_weights_in_bert_generation(
						    lowercase__      ,						lowercase__      ,						model_class='bert'      ,						is_encoder_named_decoder=lowercase__      ,						is_encoder=lowercase__      ,						)

						# Save pytorch-model
						print(f"Save PyTorch model and config to {pytorch_dump_path}"				)
						model.save_pretrained(lowercase__				)


if __name__ == "__main__":
			snake_case__     							= argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '--tf_hub_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
			)
			parser.add_argument(
			    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
			)
			parser.add_argument(
			    '--is_encoder_named_decoder',
			    action='store_true',
			    help='If decoder has to be renamed to encoder in PyTorch model.',
			)
			parser.add_argument('--is_encoder', action='store_true', help='If model is an encoder.')
			parser.add_argument('--vocab_size', default=5_0358, type=int, help='Vocab size of model')
			snake_case__     							= parser.parse_args()
			convert_tf_checkpoint_to_pytorch(
			    args.tf_hub_path,
			    args.pytorch_dump_path,
			    args.is_encoder_named_decoder,
			    args.vocab_size,
			    is_encoder=args.is_encoder,
			)


