





"""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   __lowerCamelCase     (							a_  :	List[str]							)    ->   Optional[int]:
      __SCREAMING_SNAKE_CASE     :Tuple	       =						'''https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'''

      with open(a_				,     '''wb'''							) as f:
            with requests.get(a_				,     stream=a_							) as req:
                  for chunk in req.iter_content(chunk_size=20_48							):
                        f.write(a_							)



def   __lowerCamelCase     (							)    ->   VivitConfig:
      __SCREAMING_SNAKE_CASE     :Optional[int]	       =						VivitConfig()

      __SCREAMING_SNAKE_CASE     :Optional[int]	       =						4_00
      __SCREAMING_SNAKE_CASE     :int	       =						'''huggingface/label-files'''
      __SCREAMING_SNAKE_CASE     :Tuple	       =						'''kinetics400-id2label.json'''

      __SCREAMING_SNAKE_CASE     :Any	       =						json.load(open(hf_hub_download(a_				,     a_				,     repo_type='''dataset'''							)				,     '''r'''							)							)
      __SCREAMING_SNAKE_CASE     :Any	       =						{int(a_							): v for k, v in idalabel.items()}
      __SCREAMING_SNAKE_CASE     :Tuple	       =						idalabel
      __SCREAMING_SNAKE_CASE     :List[Any]	       =						{v: k for k, v in idalabel.items()}
      return config



def   __lowerCamelCase     (							)    ->   int:
      __SCREAMING_SNAKE_CASE     :str	       =						hf_hub_download(
          repo_id='''hf-internal-testing/spaghetti-video'''				,     filename='''eating_spaghetti_32_frames.npy'''				,     repo_type='''dataset'''							)
      __SCREAMING_SNAKE_CASE     :Tuple	       =						np.load(a_							)
      return list(a_							)



def   __lowerCamelCase     (							a_  :	np.ndarray							)    ->   List[str]:
      if np.ndim(a_							) == 2:
            return transform_attention_bias(a_							)

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

      else:
            raise Exception(f'''Invalid number of dimesions: {np.ndim(a_							)}'''							)



def   __lowerCamelCase     (							a_  :	np.ndarray							)    ->   Dict:
      return current.flatten()



def   __lowerCamelCase     (							a_  :	np.ndarray							)    ->   List[str]:
      return np.reshape(a_				,     (current.shape[0], current.shape[1] * current.shape[2])							).T



def   __lowerCamelCase     (							a_  :	np.ndarray							)    ->   List[Any]:
      return np.reshape(a_				,     (current.shape[0] * current.shape[1], current.shape[2])							).T



def   __lowerCamelCase     (							a_  :	str				,     a_  :	Optional[Any]							)    ->   Optional[int]:
      __SCREAMING_SNAKE_CASE     :List[str]	       =						state_dict['''optimizer''']['''target''']['''Transformer'''][f'''encoderblock_{i}''']

      __SCREAMING_SNAKE_CASE     :str	       =						f'''encoder.layer.{i}.'''
      __SCREAMING_SNAKE_CASE     :Optional[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   __lowerCamelCase     (							a_  :	List[Any]							)    ->   Any:
      return sum([1 if '''encoderblock_''' in k else 0 for k in state_dict['''optimizer''']['''target''']['''Transformer'''].keys()]							)



def   __lowerCamelCase     (							a_  :	Tuple				,     a_  :	Optional[Any]=False							)    ->   Dict:
      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						get_n_layers(a_							)

      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						OrderedDict()

      __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''bias''']
      __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''scale''']

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

      __SCREAMING_SNAKE_CASE     :List[str]	       =						state_dict['''optimizer''']['''target''']['''cls''']
      __SCREAMING_SNAKE_CASE     :str	       =						state_dict['''optimizer''']['''target''']['''Transformer''']['''posembed_input'''][
          '''pos_embedding'''
      ]

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

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

      return {k: torch.tensor(a_							) for k, v in new_state.items()}



def   __lowerCamelCase     (							)    ->   VivitImageProcessor:
      __SCREAMING_SNAKE_CASE     :str	       =						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   __lowerCamelCase     (							a_  :	str							)    ->   str:
      __SCREAMING_SNAKE_CASE     :List[str]	       =						'''checkpoint'''

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

      __SCREAMING_SNAKE_CASE     :str	       =						restore_checkpoint(a_				,     a_							)
      __SCREAMING_SNAKE_CASE     :Optional[int]	       =						transform_state(a_				,     classification_head=a_							)

      __SCREAMING_SNAKE_CASE     :str	       =						get_vivit_config()

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

      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						VivitForVideoClassification(a_							)
      model.load_state_dict(a_							)
      model.eval()

      __SCREAMING_SNAKE_CASE     :str	       =						get_processor()

      __SCREAMING_SNAKE_CASE     :List[Any]	       =						prepare_video()
      __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						extractor(a_				,     return_tensors='''pt'''							)

      __SCREAMING_SNAKE_CASE     :List[Any]	       =						model(**a_							)

      __SCREAMING_SNAKE_CASE     :int	       =						torch.Size([1, 4_00]							)
      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]							)

      assert outputs.logits.shape == expected_shape
      assert torch.allclose(outputs.logits[0, :5]				,     a_				,     atol=1e-4							), outputs.logits[0, :5]

      model.save_pretrained(a_							)
      extractor.save_pretrained(a_							)


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

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

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






"""simple docstring"""
import math
from collections.abc import Generator



def   __lowerCamelCase     (							a_  :	int							)    ->   Generator[int, None, None]:
      __SCREAMING_SNAKE_CASE     :Generator	       =						(i for i in range(1				,     (max_n + 1)							))
      for i in (n for n in numbers if n > 1):
            for j in range(2				,     a_							):
                  if (i % j) == 0:
                        break
            else:
                  yield i



def   __lowerCamelCase     (							a_  :	int							)    ->   Generator[int, None, None]:
      __SCREAMING_SNAKE_CASE     :Generator	       =						(i for i in range(1				,     (max_n + 1)							))
      for i in (n for n in numbers if n > 1):
            # only need to check for factors up to sqrt(i)
            __SCREAMING_SNAKE_CASE     :Tuple	       =						int(math.sqrt(a_							)							) + 1
            for j in range(2				,     a_							):
                  if (i % j) == 0:
                        break
            else:
                  yield i



def   __lowerCamelCase     (							a_  :	int							)    ->   Generator[int, None, None]:
      __SCREAMING_SNAKE_CASE     :Generator	       =						(i for i in range(1				,     (max_n + 1)				,     2							))
      # It's useless to test even numbers as they will not be prime
      if max_n > 2:
            yield 2  # Because 2 will not be tested, it's necessary to yield it now
      for i in (n for n in numbers if n > 1):
            __SCREAMING_SNAKE_CASE     :int	       =						int(math.sqrt(a_							)							) + 1
            for j in range(3				,     a_				,     2							):
                  # As we removed the even numbers, we don't need them now
                  if (i % j) == 0:
                        break
            else:
                  yield i



def   __lowerCamelCase     (							)    ->   Union[str, Any]:
      from timeit import timeit

      __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						'''from __main__ import slow_primes, primes, fast_primes'''
      print(timeit('''slow_primes(1_000_000_000_000)'''				,     setup=a_				,     number=1_00_00_00							)							)
      print(timeit('''primes(1_000_000_000_000)'''				,     setup=a_				,     number=1_00_00_00							)							)
      print(timeit('''fast_primes(1_000_000_000_000)'''				,     setup=a_				,     number=1_00_00_00							)							)


if __name__ == "__main__":
  lowerCamelCase_          =	int(input("Calculate primes up to:\n>> ").strip())
  for ret in primes(number):
    print(ret)
  benchmark()