




"""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 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Optional[Any]:
	A__ =       "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=20_48		):
				f.write(lowercase_		)







def 						_SCREAMING_SNAKE_CASE     (							)		->	VivitConfig:
	A__ =       VivitConfig()

	A__ =       4_00
	A__ =       "huggingface/label-files"
	A__ =       "kinetics400-id2label.json"

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







def 						_SCREAMING_SNAKE_CASE     (							)		->	Any:
	A__ =       hf_hub_download(
	    repo_id="hf-internal-testing/spaghetti-video"    ,				filename="eating_spaghetti_32_frames.npy"    ,				repo_type="dataset"		)
	A__ =       np.load(lowercase_		)
	return list(lowercase_		)







def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Any:



	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_		)		->	Any:
	return current.flatten()







def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Union[str, Any]:
	return np.reshape(lowercase_    ,				(current.shape[0], current.shape[1] * current.shape[2])		).T







def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Optional[int]:
	return np.reshape(lowercase_    ,				(current.shape[0] * current.shape[1], current.shape[2])		).T







def 						_SCREAMING_SNAKE_CASE     (							lowercase_    ,				lowercase_		)		->	int:
	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 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	List[Any]:
	return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()]		)







def 						_SCREAMING_SNAKE_CASE     (							lowercase_    ,				lowercase_=False		)		->	Any:
	A__ =       get_n_layers(lowercase_		)

	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(lowercase_		):
		new_state.update(transform_state_encoder_block(lowercase_    ,				lowercase_		)		)

	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(lowercase_		) for k, v in new_state.items()}







def 						_SCREAMING_SNAKE_CASE     (							)		->	VivitImageProcessor:
	A__ =       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 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	List[Any]:
	A__ =       "checkpoint"

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

	A__ =       restore_checkpoint(lowercase_    ,				lowercase_		)
	A__ =       transform_state(lowercase_    ,				classification_head=lowercase_		)

	A__ =       get_vivit_config()

	assert config.image_size == 2_24
	assert config.num_frames == 32

	A__ =       VivitForVideoClassification(lowercase_		)
	model.load_state_dict(lowercase_		)
	model.eval()

	A__ =       get_processor()

	A__ =       prepare_video()
	A__ =       extractor(lowercase_    ,				return_tensors="pt"		)

	A__ =       model(**lowercase_		)

	A__ =       torch.Size([1, 4_00]		)
	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]    ,				lowercase_    ,				atol=1E-4		), outputs.logits[0, :5]

	model.save_pretrained(lowercase_		)
	extractor.save_pretrained(lowercase_		)


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

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

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











"""simple docstring"""



from typing import Optional, Union

from ...file_utils import TensorType
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, TruncationStrategy




class 			UpperCAmelCase_						(		A_							):
    lowercase__		      =   '''MarkupLMFeatureExtractor'''
    lowercase__		      =   ('''MarkupLMTokenizer''', '''MarkupLMTokenizerFast''')
    lowercase__		      =   True




    def __call__(					self		: Dict			,				snake_case_		: Optional[int]=None			,				snake_case_		: Union[str, Any]=None			,				snake_case_		: Tuple=None			,				snake_case_		: str=None			,				snake_case_		: Optional[Any]=None			,				snake_case_		: bool = True			,				snake_case_		: Union[bool, str, PaddingStrategy] = False			,				snake_case_		: Union[bool, str, TruncationStrategy] = None			,				snake_case_		: Optional[int] = None			,				snake_case_		: int = 0			,				snake_case_		: Optional[int] = None			,				snake_case_		: Optional[bool] = None			,				snake_case_		: Optional[bool] = None			,				snake_case_		: bool = False			,				snake_case_		: bool = False			,				snake_case_		: bool = False			,				snake_case_		: bool = False			,				snake_case_		: bool = True			,				snake_case_		: Optional[Union[str, TensorType]] = None			,				**snake_case_		: Dict			,				)   ->							BatchEncoding:




     '''simple docstring'''

     if self.parse_html:
      if html_strings is None:
       raise ValueError("Make sure to pass HTML strings in case `parse_html` is set to `True`"   )

      if nodes is not None or xpaths is not None or node_labels is not None:
       raise ValueError(
           "Please don't pass nodes, xpaths nor node labels in case `parse_html` is set to `True`"   )

      A__ =       self.feature_extractor(snake_case_   )
      A__ =       features["nodes"]
      A__ =       features["xpaths"]
     else:
      if html_strings is not None:
       raise ValueError("You have passed HTML strings but `parse_html` is set to `False`."   )
      if nodes is None or xpaths is None:
       raise ValueError("Make sure to pass nodes and xpaths in case `parse_html` is set to `False`"   )

        # # second, apply the tokenizer
     if questions is not None and self.parse_html:
      if isinstance(snake_case_			,				snake_case_   ):
       A__ =       [questions]  # add batch dimension (as the feature extractor always adds a batch dimension)

     A__ =       self.tokenizer(
         text=questions if questions is not None else nodes			,				text_pair=nodes if questions is not None else None			,				xpaths=snake_case_			,				node_labels=snake_case_			,				add_special_tokens=snake_case_			,				padding=snake_case_			,				truncation=snake_case_			,				max_length=snake_case_			,				stride=snake_case_			,				pad_to_multiple_of=snake_case_			,				return_token_type_ids=snake_case_			,				return_attention_mask=snake_case_			,				return_overflowing_tokens=snake_case_			,				return_special_tokens_mask=snake_case_			,				return_offsets_mapping=snake_case_			,				return_length=snake_case_			,				verbose=snake_case_			,				return_tensors=snake_case_			,				**snake_case_			,				)

     return encoded_inputs




    def 							__magic_name__							(					self		: Optional[Any]			,				*snake_case_		: Any			,				**snake_case_		: Any   )   ->							int:




     '''simple docstring'''

     return self.tokenizer.batch_decode(*snake_case_			,				**snake_case_   )




    def 							__magic_name__							(					self		: Tuple			,				*snake_case_		: Optional[Any]			,				**snake_case_		: List[Any]   )   ->							Any:




     '''simple docstring'''

     return self.tokenizer.decode(*snake_case_			,				**snake_case_   )




    @property
    def 							__magic_name__							(					self		: List[Any]   )   ->							Optional[int]:




     '''simple docstring'''

     A__ =       self.tokenizer.model_input_names
     return tokenizer_input_names





