code: 
  '''simple docstring'''
  
  
  
  
  
  import copy
  import inspect
  import unittest
  
  import numpy as np
  from huggingface_hub import hf_hub_download
  
  from transformers import VivitConfig
  from transformers.models.auto import get_values
  from transformers.testing_utils import require_torch, require_vision, slow, torch_device
  from transformers.utils import cached_property, is_torch_available, is_vision_available
  
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
  
  
  if is_torch_available():
       import torch
       from torch import nn
  
       from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VivitForVideoClassification, VivitModel
       from transformers.models.vivit.modeling_vivit import VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST
  
  
  if is_vision_available():
       from transformers import VivitImageProcessor
  
  
  
  class  lowerCAmelCase__   :
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    def __init__(self  : int	, lowerCAmelCase_  : Any	, lowerCAmelCase_  : Optional[int]=2	, lowerCAmelCase_  : Tuple=True	, lowerCAmelCase_  : List[str]=True	, lowerCAmelCase_  : Union[str, Any]=10	, lowerCAmelCase_  : Any=10	, lowerCAmelCase_  : int=8	, lowerCAmelCase_  : Optional[int]=[2, 4, 4]	, lowerCAmelCase_  : List[str]=3	, lowerCAmelCase_  : List[str]=7_68	, lowerCAmelCase_  : Union[str, Any]=5	, lowerCAmelCase_  : Any=4	, lowerCAmelCase_  : Dict=37	, lowerCAmelCase_  : Union[str, Any]="gelu_fast"	, lowerCAmelCase_  : List[Any]=0.0	, lowerCAmelCase_  : int=0.0	, lowerCAmelCase_  : List[str]=0.02	, lowerCAmelCase_  : Optional[Any]=1E-06	, lowerCAmelCase_  : List[Any]=True	, lowerCAmelCase_  : List[str]=None	, )   ->       List[Any]:
          snake_case_						       =							parent
          snake_case_						       =							batch_size
          snake_case_						       =							is_training
          snake_case_						       =							use_labels
          snake_case_						       =							num_labels
          snake_case_						       =							image_size
          snake_case_						       =							num_frames
          snake_case_						       =							tubelet_size
          snake_case_						       =							num_channels
          snake_case_						       =							hidden_size
          snake_case_						       =							num_hidden_layers
          snake_case_						       =							num_attention_heads
          snake_case_						       =							intermediate_size
          snake_case_						       =							hidden_act
          snake_case_						       =							hidden_dropout_prob
          snake_case_						       =							attention_probs_dropout_prob
          snake_case_						       =							initializer_range
          snake_case_						       =							layer_norm_eps
          snake_case_						       =							qkv_bias
          snake_case_						       =							scope
  
          snake_case_						       =							(
              (self.image_size // self.tubelet_size[2])
              * (self.image_size // self.tubelet_size[1])
              * (self.num_frames // self.tubelet_size[0])
          ) + 1  # CLS token
  
    def 		_lowerCamelCase	(self  : str    )   ->       int:
          snake_case_						       =							floats_tensor(
              [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]    )
  
          snake_case_						       =							None
          if self.use_labels:
                snake_case_						       =							ids_tensor([self.batch_size]	, self.num_labels    )
  
          snake_case_						       =							self.get_config()
  
          return config, pixel_values, labels
  
    def 		_lowerCamelCase	(self  : Dict    )   ->       Any:
          snake_case_						       =							VivitConfig(
              num_frames=self.num_frames	, image_size=self.image_size	, tubelet_size=self.tubelet_size	, num_channels=self.num_channels	, hidden_size=self.hidden_size	, num_hidden_layers=self.num_hidden_layers	, num_attention_heads=self.num_attention_heads	, intermediate_size=self.intermediate_size	, hidden_act=self.hidden_act	, hidden_dropout_prob=self.hidden_dropout_prob	, attention_probs_dropout_prob=self.attention_probs_dropout_prob	, initializer_range=self.initializer_range	, layer_norm_eps=self.layer_norm_eps	, qkv_bias=self.qkv_bias	, )
          snake_case_						       =							self.num_labels
          return config
  
    def 		_lowerCamelCase	(self  : Tuple	, lowerCAmelCase_  : List[Any]	, lowerCAmelCase_  : List[str]	, lowerCAmelCase_  : Optional[int]    )   ->       Dict:
          snake_case_						       =							VivitModel(config=lowerCAmelCase_    )
          model.to(lowerCAmelCase_    )
          model.eval()
          snake_case_						       =							model(lowerCAmelCase_    )
          self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, self.seq_length, self.hidden_size)    )
  
    def 		_lowerCamelCase	(self  : List[Any]	, lowerCAmelCase_  : Any	, lowerCAmelCase_  : Union[str, Any]	, lowerCAmelCase_  : str    )   ->       str:
          snake_case_						       =							VivitForVideoClassification(lowerCAmelCase_    )
          model.to(lowerCAmelCase_    )
          model.eval()
  
          snake_case_						       =							model(lowerCAmelCase_    )
  
          # verify the logits shape
          snake_case_						       =							torch.Size((self.batch_size, self.num_labels)    )
          self.parent.assertEqual(result.logits.shape	, lowerCAmelCase_    )
  
  
    def 		_lowerCamelCase	(self  : Optional[int]    )   ->       int:
          snake_case_						       =							self.prepare_config_and_inputs()
          snake_case_       ,	snake_case_       ,	snake_case_						       =							config_and_inputs
          snake_case_						       =							{"pixel_values": pixel_values}
          return config, inputs_dict
  
  
  
  @require_torch
  class  lowerCAmelCase__   (       UpperCamelCase					,							unittest.TestCase       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    _UpperCamelCase						= (VivitModel, VivitForVideoClassification) if is_torch_available() else ()
  
    _UpperCamelCase						= False
    _UpperCamelCase						= False
    _UpperCamelCase						= False
    _UpperCamelCase						= False
  
    def 		_lowerCamelCase	(self  : str    )   ->       Any:
          snake_case_						       =							VivitModelTester(self    )
          snake_case_						       =							ConfigTester(self	, config_class=lowerCAmelCase_	, has_text_modality=lowerCAmelCase_	, hidden_size=37    )
  
    def 		_lowerCamelCase	(self  : Any	, lowerCAmelCase_  : List[str]	, lowerCAmelCase_  : str	, lowerCAmelCase_  : Union[str, Any]=False    )   ->       Optional[int]:
          snake_case_						       =							copy.deepcopy(lowerCAmelCase_    )
  
          if return_labels:
                if model_class in get_values(lowerCAmelCase_    ):
                      snake_case_						       =							torch.zeros(
                          self.model_tester.batch_size	, dtype=torch.long	, device=lowerCAmelCase_    )
  
          return inputs_dict
  
    def 		_lowerCamelCase	(self  : Optional[Any]    )   ->       Optional[int]:
          self.config_tester.run_common_tests()
  
    @unittest.skip(reason="Vivit does not use inputs_embeds"    )
    def 		_lowerCamelCase	(self  : Tuple    )   ->       Optional[int]:
          pass
  
    def 		_lowerCamelCase	(self  : int    )   ->       Optional[Any]:
          snake_case_       ,	snake_case_						       =							self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
                snake_case_						       =							model_class(lowerCAmelCase_    )
                self.assertIsInstance(model.get_input_embeddings()	, (nn.Module)    )
                snake_case_						       =							model.get_output_embeddings()
                self.assertTrue(x is None or isinstance(lowerCAmelCase_	, nn.Linear    )    )
  
    def 		_lowerCamelCase	(self  : Any    )   ->       List[str]:
          snake_case_       ,	snake_case_						       =							self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
                snake_case_						       =							model_class(lowerCAmelCase_    )
                snake_case_						       =							inspect.signature(model.forward    )
                # signature.parameters is an OrderedDict => so arg_names order is deterministic
                snake_case_						       =							[*signature.parameters.keys()]
  
                snake_case_						       =							["pixel_values", "head_mask"]
                self.assertListEqual(arg_names[:2]	, lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : Any    )   ->       Tuple:
          snake_case_						       =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : int    )   ->       Optional[int]:
          snake_case_						       =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_video_classification(*lowerCAmelCase_    )
  
    @slow
    def 		_lowerCamelCase	(self  : Dict    )   ->       Union[str, Any]:
          for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                snake_case_						       =							VivitModel.from_pretrained(lowerCAmelCase_    )
                self.assertIsNotNone(lowerCAmelCase_    )
  
    def 		_lowerCamelCase	(self  : Dict    )   ->       Dict:
          snake_case_       ,	snake_case_						       =							self.model_tester.prepare_config_and_inputs_for_common()
          snake_case_						       =							True
  
          for model_class in self.all_model_classes:
                snake_case_						       =							self.model_tester.seq_length
  
                snake_case_						       =							True
                snake_case_						       =							False
                snake_case_						       =							True
                snake_case_						       =							model_class(lowerCAmelCase_    )
                model.to(lowerCAmelCase_    )
                model.eval()
                with torch.no_grad():
                      snake_case_						       =							model(**self._prepare_for_class(lowerCAmelCase_	, lowerCAmelCase_    )    )
                snake_case_						       =							outputs.attentions
                self.assertEqual(len(lowerCAmelCase_    )	, self.model_tester.num_hidden_layers    )
  
                # check that output_attentions also work using config
                del inputs_dict["output_attentions"]
                snake_case_						       =							True
                snake_case_						       =							model_class(lowerCAmelCase_    )
                model.to(lowerCAmelCase_    )
                model.eval()
                with torch.no_grad():
                      snake_case_						       =							model(**self._prepare_for_class(lowerCAmelCase_	, lowerCAmelCase_    )    )
                snake_case_						       =							outputs.attentions
                self.assertEqual(len(lowerCAmelCase_    )	, self.model_tester.num_hidden_layers    )
  
                self.assertListEqual(
                    list(attentions[0].shape[-3:]    )	, [self.model_tester.num_attention_heads, seq_len, seq_len]	, )
                snake_case_						       =							len(lowerCAmelCase_    )
  
                # Check attention is always last and order is fine
                snake_case_						       =							True
                snake_case_						       =							True
                snake_case_						       =							model_class(lowerCAmelCase_    )
                model.to(lowerCAmelCase_    )
                model.eval()
                with torch.no_grad():
                      snake_case_						       =							model(**self._prepare_for_class(lowerCAmelCase_	, lowerCAmelCase_    )    )
  
                self.assertEqual(out_len + 1	, len(lowerCAmelCase_    )    )
  
                snake_case_						       =							outputs.attentions
  
                self.assertEqual(len(lowerCAmelCase_    )	, self.model_tester.num_hidden_layers    )
  
                self.assertListEqual(
                    list(self_attentions[0].shape[-3:]    )	, [self.model_tester.num_attention_heads, seq_len, seq_len]	, )
  
    def 		_lowerCamelCase	(self  : Optional[int]    )   ->       Optional[int]:
          def check_hidden_states_output(lowerCAmelCase_  : int	, lowerCAmelCase_  : Union[str, Any]	, lowerCAmelCase_  : List[str]    ):
                snake_case_						       =							model_class(lowerCAmelCase_    )
                model.to(lowerCAmelCase_    )
                model.eval()
  
                with torch.no_grad():
                      snake_case_						       =							model(**self._prepare_for_class(lowerCAmelCase_	, lowerCAmelCase_    )    )
  
                snake_case_						       =							outputs.hidden_states
                snake_case_						       =							self.model_tester.num_hidden_layers + 1
                self.assertEqual(len(lowerCAmelCase_    )	, lowerCAmelCase_    )
  
                snake_case_						       =							self.model_tester.seq_length
  
                self.assertListEqual(
                    list(hidden_states[0].shape[-2:]    )	, [seq_length, self.model_tester.hidden_size]	, )
  
          snake_case_       ,	snake_case_						       =							self.model_tester.prepare_config_and_inputs_for_common()
  
          for model_class in self.all_model_classes:
                snake_case_						       =							True
                check_hidden_states_output(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
  
                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                snake_case_						       =							True
  
                check_hidden_states_output(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
  
  
    @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."    )
    def 		_lowerCamelCase	(self  : Union[str, Any]    )   ->       str:
          pass
  
  def   snake_case_       (							):
        snake_case_						       =							hf_hub_download(
            repo_id="hf-internal-testing/spaghetti-video" ,					filename="eating_spaghetti_32_frames.npy" ,					repo_type="dataset"			)
        snake_case_						       =							np.load(_a			)
        return list(_a			)
  
  
  
  
  @require_torch
  @require_vision
  class  lowerCAmelCase__   (       unittest.TestCase       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    @cached_property
    def 		_lowerCamelCase	(self  : Dict    )   ->       Tuple:
          return VivitImageProcessor() if is_vision_available() else None
  
  
    @slow
    def 		_lowerCamelCase	(self  : Union[str, Any]    )   ->       List[str]:
          snake_case_						       =							VivitForVideoClassification.from_pretrained("google/vivit-b-16x2-kinetics400"    ).to(lowerCAmelCase_    )
  
          snake_case_						       =							self.default_image_processor
          snake_case_						       =							prepare_video()
          snake_case_						       =							image_processor(lowerCAmelCase_	, return_tensors="pt"    ).to(lowerCAmelCase_    )
  
          # forward pass
          with torch.no_grad():
                snake_case_						       =							model(**lowerCAmelCase_    )
  
          # verify the logits
          snake_case_						       =							torch.Size((1, 4_00)    )
          self.assertEqual(outputs.logits.shape	, lowerCAmelCase_    )
  
          # taken from original model
          snake_case_						       =							torch.tensor([-1.0_543, 2.0_764, -0.2_104, 0.4_439, -0.9_658]    ).to(lowerCAmelCase_    )
  
          self.assertTrue(torch.allclose(outputs.logits[0, :5]	, lowerCAmelCase_	, atol=1E-4    )    )
  
  
style_context: 
  '''simple docstring'''
  
  
  
  
  
  import jax
  import jax.numpy as jnp
  from bigbird_flax import FlaxBigBirdForNaturalQuestions
  from datasets import load_from_disk
  
  from transformers import BigBirdTokenizerFast
  
  
  A           =   {0: "null", 1: "short", 2: "long", 3: "yes", 4: "no"}
  A           =   set("".join(["‘", "’", "´", "`", ".", ",", "-", "\""]))
  def   snake_case_       (							_a     :							Any ,					_a     :							Union[str, Any]=0 ,					_a     :							Optional[int]=None			):
        return [" ".join(x.split(" "			)[begin:end]			) for x in answers if len(x.split(" "			)			) > 1]
  def   snake_case_       (							_a     :							Optional[int] ,					_a     :							str=False			):
        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"
              snake_case_						       =							(
                  given_answers + get_sub_answers(_a ,					begin=1			) + get_sub_answers(_a ,					end=-1			)
              )
        snake_case_						       =							[]
        for answer in given_answers:
              snake_case_						       =							answer.replace("_" ,					" "			).lower()
              snake_case_						       =							"".join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else " " for c in alias			)
              answers.append(" ".join(alias.split()			).strip()			)
        return set(_a			)
  def   snake_case_       (							_a     :							Optional[Any] ,					_a     :							int ,					_a     :							Optional[Any]=1 ,					_a     :							List[Any]=100			):
        snake_case_       ,	snake_case_						       =							jax.lax.top_k(_a ,					_a			)
        snake_case_       ,	snake_case_						       =							jax.lax.top_k(_a ,					_a			)
  
        snake_case_						       =							best_end_idx[:, None] - best_start_idx[None, :]
        snake_case_						       =							jnp.logical_or(widths < 0 ,					widths > max_size			)
        snake_case_						       =							(best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
        snake_case_						       =							jnp.argmax(_a			).item()
  
        return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]
  def   snake_case_       (							_a     :							Any			):
        snake_case_						       =							sample["question"]["text"]
        snake_case_						       =							sample["document"]["tokens"]["token"]
        snake_case_						       =							sample["document"]["tokens"]["is_html"]
        snake_case_						       =							sample["annotations"]["long_answer"]
        snake_case_						       =							sample["annotations"]["short_answers"]
  
        snake_case_						       =							" ".join([context[i] for i in range(len(_a			)			) 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",
                    }
  
        snake_case_						       =							[]
        for s in short_answers:
              short_targets.extend(s["text"]			)
        snake_case_						       =							list(set(_a			)			)
  
        snake_case_						       =							[]
        for s in long_answers:
              if s["start_token"] == -1:
                    continue
              snake_case_						       =							context[s["start_token"] : s["end_token"]]
              snake_case_						       =							is_html[s["start_token"] : s["end_token"]]
              snake_case_						       =							" ".join([answer[i] for i in range(len(_a			)			) if not html[i]]			)
              if new_answer not in long_targets:
                    long_targets.append(_a			)
  
        snake_case_						       =							"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   snake_case_       (							):
        snake_case_						       =							load_from_disk("natural-questions-validation"			)
        snake_case_						       =							dataset.map(_a			).remove_columns(["annotations", "document", "id"]			)
        print(_a			)
  
        snake_case_						       =							dataset.filter(lambda _a			: (len(x["question"]			) + len(x["context"]			)) < 4 * 4_096			)
        snake_case_						       =							short_validation_dataset.filter(lambda _a			: x["category"] != "null"			)
        short_validation_dataset
  
        snake_case_						       =							"vasudevgupta/flax-bigbird-natural-questions"
        snake_case_						       =							FlaxBigBirdForNaturalQuestions.from_pretrained(_a			)
        snake_case_						       =							BigBirdTokenizerFast.from_pretrained(_a			)
  
        @jax.jit
        def forward(*_a     :							List[Any] ,					**_a     :							int			):
              snake_case_       ,	snake_case_       ,	snake_case_						       =							model(*_a ,					**_a			)
              return start_logits, end_logits, jnp.argmax(_a ,					axis=-1			)
  
        def evaluate(_a     :							List[Any]			):
              # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
              snake_case_						       =							tokenizer(
                  example["question"] ,					example["context"] ,					return_tensors="np" ,					max_length=4_096 ,					padding="max_length" ,					truncation=_a ,					)
  
              snake_case_       ,	snake_case_       ,	snake_case_						       =							forward(**_a			)
  
              snake_case_						       =							CATEGORY_MAPPING[category.item()]
  
              snake_case_						       =							example["long"] + example["short"]
              if example["category"] in ["yes", "no", "null"]:
                    snake_case_						       =							[example["category"]]
              snake_case_						       =							example["category"] != "null"
              # Now target can be: "yes", "no", "null", "list of long & short answers"
  
              if predicted_category in ["yes", "no", "null"]:
                    snake_case_						       =							[predicted_category]
                    snake_case_						       =							example["output"] == example["targets"]
                    snake_case_						       =							predicted_category != "null"
                    return example
  
              snake_case_						       =							38 if predicted_category == "short" else 1_024
              snake_case_       ,	snake_case_						       =							get_best_valid_start_end_idx(
                  start_scores[0] ,					end_scores[0] ,					top_k=8 ,					max_size=_a			)
  
              snake_case_						       =							inputs["input_ids"][0].tolist()
              snake_case_						       =							[tokenizer.decode(input_ids[start_score : end_score + 1]			)]
  
              snake_case_						       =							expand_to_aliases(example["targets"] ,					make_sub_answers=_a			)
              snake_case_						       =							expand_to_aliases(example["output"]			)
  
              # some preprocessing to both prediction and answer
              snake_case_						       =							{"".join(a.split()			) for a in answers}
              snake_case_						       =							{"".join(p.split()			) for p in predictions}
              snake_case_						       =							{s for s in predictions if s not in ["``", "''", "`", "'"]}
  
              # if there is a common element, it's a exact match
              snake_case_						       =							len(list(answers & predictions			)			) > 0
              snake_case_						       =							predicted_category != "null" and len(_a			) > 0
  
              return example
  
        snake_case_						       =							short_validation_dataset.map(_a			)
  
        snake_case_						       =							len(_a			)
        snake_case_						       =							len(short_validation_dataset.filter(lambda _a			: x["match"] == 1			)			)
        print("EM score:" ,					(matched / total) * 100 ,					"%"			)
  
  
  if __name__ == "__main__":
       main()
  
  