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 		_A		(		_lowercase   )				->   int:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = '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 		_A		(		)				->   VivitConfig:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = VivitConfig()
  
       __UpperCamelCase   = 4_00
       __UpperCamelCase   = 'huggingface/label-files'
       __UpperCamelCase   = 'kinetics400-id2label.json'
  
       __UpperCamelCase   = json.load(open(hf_hub_download(_lowercase				,     _lowercase				,     repo_type='dataset'   )				,     'r'   )   )
       __UpperCamelCase   = {int(_lowercase   ): v for k, v in idalabel.items()}
       __UpperCamelCase   = idalabel
       __UpperCamelCase   = {v: k for k, v in idalabel.items()}
       return config
  
  
  
  
  
  
  def 		_A		(		)				->   Optional[Any]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = hf_hub_download(
           repo_id='hf-internal-testing/spaghetti-video'				,     filename='eating_spaghetti_32_frames.npy'				,     repo_type='dataset'   )
       __UpperCamelCase   = np.load(_lowercase   )
       return list(_lowercase   )
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   List[str]:
  
  
       """simple docstring"""
  
  
       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 		_A		(		_lowercase   )				->   str:
  
  
       """simple docstring"""
  
  
       return current.flatten()
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   Dict:
  
  
       """simple docstring"""
  
  
       return np.reshape(_lowercase				,     (current.shape[0], current.shape[1] * current.shape[2])   ).T
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   Optional[Any]:
  
  
       """simple docstring"""
  
  
       return np.reshape(_lowercase				,     (current.shape[0] * current.shape[1], current.shape[2])   ).T
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase   )				->   Any:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = state_dict['optimizer']['target']['Transformer'][f'''encoderblock_{i}''']
  
       __UpperCamelCase   = f'''encoder.layer.{i}.'''
       __UpperCamelCase   = {
           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		(		_lowercase   )				->   Optional[int]:
  
  
       """simple docstring"""
  
  
       return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]   )
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase=False   )				->   List[str]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = get_n_layers(_lowercase   )
  
       __UpperCamelCase   = OrderedDict()
  
       __UpperCamelCase   = state_dict['optimizer']['target']['Transformer']['encoder_norm']['bias']
       __UpperCamelCase   = state_dict['optimizer']['target']['Transformer']['encoder_norm']['scale']
  
       __UpperCamelCase   = np.transpose(
           state_dict['optimizer']['target']['embedding']['kernel']				,     (4, 3, 0, 1, 2)   )
       __UpperCamelCase   = state_dict['optimizer']['target']['embedding']['bias']
  
       __UpperCamelCase   = state_dict['optimizer']['target']['cls']
       __UpperCamelCase   = 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:
            __UpperCamelCase   = {'vivit.' + k: v for k, v in new_state.items()}
            __UpperCamelCase   = np.transpose(state_dict['optimizer']['target']['output_projection']['kernel']   )
            __UpperCamelCase   = np.transpose(state_dict['optimizer']['target']['output_projection']['bias']   )
  
       return {k: torch.tensor(_lowercase   ) for k, v in new_state.items()}
  
  
  
  
  
  
  def 		_A		(		)				->   VivitImageProcessor:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = 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 		_A		(		_lowercase   )				->   Union[str, Any]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = 'checkpoint'
  
       if not os.path.exists(_lowercase   ):
            download_checkpoint(_lowercase   )
  
       __UpperCamelCase   = restore_checkpoint(_lowercase				,     _lowercase   )
       __UpperCamelCase   = transform_state(_lowercase				,     classification_head=_lowercase   )
  
       __UpperCamelCase   = get_vivit_config()
  
       assert config.image_size == 2_24
       assert config.num_frames == 32
  
       __UpperCamelCase   = VivitForVideoClassification(_lowercase   )
       model.load_state_dict(_lowercase   )
       model.eval()
  
       __UpperCamelCase   = get_processor()
  
       __UpperCamelCase   = prepare_video()
       __UpperCamelCase   = extractor(_lowercase				,     return_tensors='pt'   )
  
       __UpperCamelCase   = model(**_lowercase   )
  
       __UpperCamelCase   = torch.Size([1, 4_00]   )
       __UpperCamelCase   = torch.tensor([-1.05_43, 2.07_64, -0.21_04, 0.44_39, -0.96_58]   )
  
       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)
  
  
  
style_context: 
  
  
  import jax
  import jax.numpy as jnp
  from bigbird_flax import FlaxBigBirdForNaturalQuestions
  from datasets import load_from_disk
  
  from transformers import BigBirdTokenizerFast
  
  
  __snake_case             =	{0: '''null''', 1: '''short''', 2: '''long''', 3: '''yes''', 4: '''no'''}
  __snake_case             =	set(''''''.join(['''‘''', '''’''', '''´''', '''`''', '''.''', ''',''', '''-''', '''"''']))
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase=0				,     _lowercase=None   )				->   Union[str, Any]:
  
  
       """simple docstring"""
  
  
       return [" ".join(x.split(' '   )[begin:end]   ) for x in answers if len(x.split(' '   )   ) > 1]
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase=False   )				->   Tuple:
  
  
       """simple docstring"""
  
  
       if make_sub_answers:
            # if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
            # *e.g.* if the correct answer contains a prefix such as "the", or "a"
            __UpperCamelCase   = (
                given_answers + get_sub_answers(_lowercase				,     begin=1   ) + get_sub_answers(_lowercase				,     end=-1   )
            )
       __UpperCamelCase   = []
       for answer in given_answers:
            __UpperCamelCase   = answer.replace('_'				,     ' '   ).lower()
            __UpperCamelCase   = ''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ' ' for c in alias   )
            answers.append(' '.join(alias.split()   ).strip()   )
       return set(_lowercase   )
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase				,     _lowercase=1				,     _lowercase=1_00   )				->   Any:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase,						__UpperCamelCase   = jax.lax.top_k(_lowercase				,     _lowercase   )
       __UpperCamelCase,						__UpperCamelCase   = jax.lax.top_k(_lowercase				,     _lowercase   )
  
       __UpperCamelCase   = best_end_idx[:, None] - best_start_idx[None, :]
       __UpperCamelCase   = jnp.logical_or(widths < 0				,     widths > max_size   )
       __UpperCamelCase   = (best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
       __UpperCamelCase   = jnp.argmax(_lowercase   ).item()
  
       return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   List[str]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = sample['question']['text']
       __UpperCamelCase   = sample['document']['tokens']['token']
       __UpperCamelCase   = sample['document']['tokens']['is_html']
       __UpperCamelCase   = sample['annotations']['long_answer']
       __UpperCamelCase   = sample['annotations']['short_answers']
  
       __UpperCamelCase   = ' '.join([context[i] for i in range(len(_lowercase   )   ) if not is_html[i]]   )
  
       # 0 - No ; 1 - Yes
       for answer in sample["annotations"]["yes_no_answer"]:
            if answer == 0 or answer == 1:
                 return {
                     "question": question,
                     "context": context_string,
                     "short": [],
                     "long": [],
                     "category": "no" if answer == 0 else "yes",
                 }
  
       __UpperCamelCase   = []
       for s in short_answers:
            short_targets.extend(s['text']   )
       __UpperCamelCase   = list(set(_lowercase   )   )
  
       __UpperCamelCase   = []
       for s in long_answers:
            if s["start_token"] == -1:
                 continue
            __UpperCamelCase   = context[s['start_token'] : s['end_token']]
            __UpperCamelCase   = is_html[s['start_token'] : s['end_token']]
            __UpperCamelCase   = ' '.join([answer[i] for i in range(len(_lowercase   )   ) if not html[i]]   )
            if new_answer not in long_targets:
                 long_targets.append(_lowercase   )
  
       __UpperCamelCase   = 'long_short' if len(short_targets + long_targets   ) > 0 else 'null'
  
       return {
           "question": question,
           "context": context_string,
           "short": short_targets,
           "long": long_targets,
           "category": category,
       }
  
  
  
  
  
  
  def 		_A		(		)				->   Optional[int]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = load_from_disk('natural-questions-validation'   )
       __UpperCamelCase   = dataset.map(_lowercase   ).remove_columns(['annotations', 'document', 'id']   )
       print(_lowercase   )
  
       __UpperCamelCase   = dataset.filter(lambda _lowercase   : (len(x['question']   ) + len(x['context']   )) < 4 * 40_96   )
       __UpperCamelCase   = short_validation_dataset.filter(lambda _lowercase   : x["category"] != "null"   )
       short_validation_dataset
  
       __UpperCamelCase   = 'vasudevgupta/flax-bigbird-natural-questions'
       __UpperCamelCase   = FlaxBigBirdForNaturalQuestions.from_pretrained(_lowercase   )
       __UpperCamelCase   = BigBirdTokenizerFast.from_pretrained(_lowercase   )
  
       @jax.jit
       def forward(*_lowercase				,     **_lowercase   ):
            __UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase   = model(*_lowercase				,     **_lowercase   )
            return start_logits, end_logits, jnp.argmax(_lowercase				,     axis=-1   )
  
       def evaluate(_lowercase   ):
            # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
            __UpperCamelCase   = tokenizer(
                example['question']				,     example['context']				,     return_tensors='np'				,     max_length=40_96				,     padding='max_length'				,     truncation=_lowercase				,     )
  
            __UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase   = forward(**_lowercase   )
  
            __UpperCamelCase   = CATEGORY_MAPPING[category.item()]
  
            __UpperCamelCase   = example['long'] + example['short']
            if example["category"] in ["yes", "no", "null"]:
                 __UpperCamelCase   = [example['category']]
            __UpperCamelCase   = example['category'] != 'null'
            # Now target can be: "yes", "no", "null", "list of long & short answers"
  
            if predicted_category in ["yes", "no", "null"]:
                 __UpperCamelCase   = [predicted_category]
                 __UpperCamelCase   = example['output'] == example['targets']
                 __UpperCamelCase   = predicted_category != 'null'
                 return example
  
            __UpperCamelCase   = 38 if predicted_category == 'short' else 10_24
            __UpperCamelCase,						__UpperCamelCase   = get_best_valid_start_end_idx(
                start_scores[0]				,     end_scores[0]				,     top_k=8				,     max_size=_lowercase   )
  
            __UpperCamelCase   = inputs['input_ids'][0].tolist()
            __UpperCamelCase   = [tokenizer.decode(input_ids[start_score : end_score + 1]   )]
  
            __UpperCamelCase   = expand_to_aliases(example['targets']				,     make_sub_answers=_lowercase   )
            __UpperCamelCase   = expand_to_aliases(example['output']   )
  
            # some preprocessing to both prediction and answer
            __UpperCamelCase   = {''.join(a.split()   ) for a in answers}
            __UpperCamelCase   = {''.join(p.split()   ) for p in predictions}
            __UpperCamelCase   = {s for s in predictions if s not in ['``', '\'\'', '`', '\'']}
  
            # if there is a common element, it's a exact match
            __UpperCamelCase   = len(list(answers & predictions   )   ) > 0
            __UpperCamelCase   = predicted_category != 'null' and len(_lowercase   ) > 0
  
            return example
  
       __UpperCamelCase   = short_validation_dataset.map(_lowercase   )
  
       __UpperCamelCase   = len(_lowercase   )
       __UpperCamelCase   = len(short_validation_dataset.filter(lambda _lowercase   : x["match"] == 1   )   )
       print('EM score:'				,     (matched / total) * 1_00				,     '%'   )
  
  
  if __name__ == "__main__":
      main()
  
  
  