





import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
      import jax
      import jax.numpy as jnp

      from transformers.generation import (
          FlaxForcedBOSTokenLogitsProcessor,
          FlaxForcedEOSTokenLogitsProcessor,
          FlaxLogitsProcessorList,
          FlaxMinLengthLogitsProcessor,
          FlaxTemperatureLogitsWarper,
          FlaxTopKLogitsWarper,
          FlaxTopPLogitsWarper,
      )






@require_flax
class        snake_case_    (		unittest.TestCase						):






  def 				UpperCAmelCase__						(							self			:					Any       ,  _snake_case			:					int       ,  _snake_case			:					int      )->str:






   '''simple docstring'''

   __lowerCAmelCase					:	Any      =       jnp.ones((batch_size, length)      ) / length
   return scores






  def 				UpperCAmelCase__						(							self			:					List[Any]      )->List[str]:






   '''simple docstring'''

   __lowerCAmelCase					:	Any      =       None
   __lowerCAmelCase					:	Optional[int]      =       20

   __lowerCAmelCase					:	Union[str, Any]      =       self._get_uniform_logits(batch_size=2       ,  length=_lowerCamelCase      )

   # tweak scores to not be uniform anymore
   __lowerCAmelCase					:	Union[str, Any]      =       scores.at[1, 5].set((1 / length) + 0.1      )  # peak, 1st batch
   __lowerCAmelCase					:	Tuple      =       scores.at[1, 10].set((1 / length) - 0.4      )  # valley, 1st batch

   # compute softmax
   __lowerCAmelCase					:	Any      =       jax.nn.softmax(_lowerCamelCase       ,  axis=-1      )

   __lowerCAmelCase					:	Dict      =       FlaxTemperatureLogitsWarper(temperature=0.5      )
   __lowerCAmelCase					:	Any      =       FlaxTemperatureLogitsWarper(temperature=1.3      )

   __lowerCAmelCase					:	Dict      =       jax.nn.softmax(temp_dist_warper_sharper(_lowerCamelCase       ,  scores.copy()       ,  cur_len=_lowerCamelCase      )       ,  axis=-1      )
   __lowerCAmelCase					:	Any      =       jax.nn.softmax(temp_dist_warper_smoother(_lowerCamelCase       ,  scores.copy()       ,  cur_len=_lowerCamelCase      )       ,  axis=-1      )

   # uniform distribution stays uniform
   self.assertTrue(jnp.allclose(probs[0, :]       ,  warped_prob_sharp[0, :]       ,  atol=1E-3      )      )
   self.assertTrue(jnp.allclose(probs[0, :]       ,  warped_prob_smooth[0, :]       ,  atol=1E-3      )      )

   # sharp peaks get higher, valleys get lower
   self.assertLess(probs[1, :].max()       ,  warped_prob_sharp[1, :].max()      )
   self.assertGreater(probs[1, :].min()       ,  warped_prob_sharp[1, :].min()      )

   # smooth peaks get lower, valleys get higher
   self.assertGreater(probs[1, :].max()       ,  warped_prob_smooth[1, :].max()      )
   self.assertLess(probs[1, :].min()       ,  warped_prob_smooth[1, :].min()      )






  def 				UpperCAmelCase__						(							self			:					Dict      )->Dict:






   '''simple docstring'''

   __lowerCAmelCase					:	Union[str, Any]      =       None
   __lowerCAmelCase					:	List[str]      =       10
   __lowerCAmelCase					:	int      =       2

   # create ramp distribution
   __lowerCAmelCase					:	int      =       np.broadcast_to(np.arange(_lowerCamelCase      )[None, :]       ,  (batch_size, vocab_size)      ).copy()
   __lowerCAmelCase					:	Tuple      =       ramp_logits[1:, : vocab_size // 2] + vocab_size

   __lowerCAmelCase					:	Optional[Any]      =       FlaxTopKLogitsWarper(3      )

   __lowerCAmelCase					:	Any      =       top_k_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )

   # check that correct tokens are filtered
   self.assertListEqual(jnp.isinf(scores[0]      ).tolist()       ,  7 * [True] + 3 * [False]      )
   self.assertListEqual(jnp.isinf(scores[1]      ).tolist()       ,  2 * [True] + 3 * [False] + 5 * [True]      )

   # check special case
   __lowerCAmelCase					:	List[str]      =       5
   __lowerCAmelCase					:	str      =       FlaxTopKLogitsWarper(top_k=1       ,  filter_value=0.0       ,  min_tokens_to_keep=3      )

   __lowerCAmelCase					:	Optional[int]      =       np.broadcast_to(np.arange(_lowerCamelCase      )[None, :]       ,  (batch_size, length)      ).copy()
   __lowerCAmelCase					:	Dict      =       top_k_warp_safety_check(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )

   # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
   self.assertListEqual((scores == 0.0).sum(axis=-1      ).tolist()       ,  [2, 2]      )






  def 				UpperCAmelCase__						(							self			:					List[Any]      )->Tuple:






   '''simple docstring'''

   __lowerCAmelCase					:	List[Any]      =       None
   __lowerCAmelCase					:	Dict      =       10
   __lowerCAmelCase					:	List[Any]      =       2

   # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
   __lowerCAmelCase					:	List[Any]      =       np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]      )      )

   __lowerCAmelCase					:	Optional[Any]      =       FlaxTopPLogitsWarper(0.8      )
   __lowerCAmelCase					:	List[Any]      =       np.exp(top_p_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )      )

   # dist should be filtered to keep min num values so that sum is >= top_p
   # exp (-inf) => 0
   __lowerCAmelCase					:	List[str]      =       np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]      )
   self.assertTrue(np.allclose(_lowerCamelCase       ,  _lowerCamelCase       ,  atol=1E-3      )      )

   # check edge cases with negative and extreme logits
   __lowerCAmelCase					:	Optional[int]      =       np.broadcast_to(np.arange(_lowerCamelCase      )[None, :]       ,  (batch_size, vocab_size)      ).copy() - (
       vocab_size // 2
   )

   # make ramp_logits more extreme
   __lowerCAmelCase					:	Dict      =       ramp_logits[1] * 100.0

   # make sure at least 2 tokens are kept
   __lowerCAmelCase					:	Tuple      =       FlaxTopPLogitsWarper(0.9       ,  min_tokens_to_keep=2       ,  filter_value=0.0      )
   __lowerCAmelCase					:	List[str]      =       top_p_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )

   # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
   self.assertListEqual((filtered_dist != 0.0).sum(axis=-1      ).tolist()       ,  [3, 2]      )






  def 				UpperCAmelCase__						(							self			:					Optional[Any]      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	Any      =       20
   __lowerCAmelCase					:	List[Any]      =       4
   __lowerCAmelCase					:	Union[str, Any]      =       0

   __lowerCAmelCase					:	Tuple      =       FlaxMinLengthLogitsProcessor(min_length=10       ,  eos_token_id=_lowerCamelCase      )

   # check that min length is applied at length 5
   __lowerCAmelCase					:	List[Any]      =       ids_tensor((batch_size, 20)       ,  vocab_size=20      )
   __lowerCAmelCase					:	List[Any]      =       5
   __lowerCAmelCase					:	Union[str, Any]      =       self._get_uniform_logits(_lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	Tuple      =       min_dist_processor(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()       ,  4 * [-float("""inf"""      )]      )

   # check that min length is not applied anymore at length 15
   __lowerCAmelCase					:	str      =       self._get_uniform_logits(_lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	int      =       15
   __lowerCAmelCase					:	Tuple      =       min_dist_processor(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   self.assertFalse(jnp.isinf(_lowerCamelCase      ).any()      )






  def 				UpperCAmelCase__						(							self			:					List[Any]      )->Dict:






   '''simple docstring'''

   __lowerCAmelCase					:	Any      =       20
   __lowerCAmelCase					:	str      =       4
   __lowerCAmelCase					:	str      =       0

   __lowerCAmelCase					:	Any      =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCamelCase      )

   # check that all scores are -inf except the bos_token_id score
   __lowerCAmelCase					:	Tuple      =       ids_tensor((batch_size, 1)       ,  vocab_size=20      )
   __lowerCAmelCase					:	Tuple      =       1
   __lowerCAmelCase					:	str      =       self._get_uniform_logits(_lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	Union[str, Any]      =       logits_processor(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]      ).all()      )
   self.assertListEqual(scores[:, bos_token_id].tolist()       ,  4 * [0]      )  # score for bos_token_id shold be zero

   # check that bos_token_id is not forced if current length is greater than 1
   __lowerCAmelCase					:	Any      =       3
   __lowerCAmelCase					:	Any      =       self._get_uniform_logits(_lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	Union[str, Any]      =       logits_processor(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   self.assertFalse(jnp.isinf(_lowerCamelCase      ).any()      )






  def 				UpperCAmelCase__						(							self			:					str      )->Tuple:






   '''simple docstring'''

   __lowerCAmelCase					:	int      =       20
   __lowerCAmelCase					:	int      =       4
   __lowerCAmelCase					:	List[Any]      =       0
   __lowerCAmelCase					:	List[str]      =       5

   __lowerCAmelCase					:	Union[str, Any]      =       FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCamelCase       ,  eos_token_id=_lowerCamelCase      )

   # check that all scores are -inf except the eos_token_id when max_length is reached
   __lowerCAmelCase					:	List[str]      =       ids_tensor((batch_size, 4)       ,  vocab_size=20      )
   __lowerCAmelCase					:	List[str]      =       4
   __lowerCAmelCase					:	int      =       self._get_uniform_logits(_lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	Optional[Any]      =       logits_processor(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]      ).all()      )
   self.assertListEqual(scores[:, eos_token_id].tolist()       ,  4 * [0]      )  # score for eos_token_id should be zero

   # check that eos_token_id is not forced if max_length is not reached
   __lowerCAmelCase					:	Optional[Any]      =       3
   __lowerCAmelCase					:	str      =       self._get_uniform_logits(_lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	str      =       logits_processor(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   self.assertFalse(jnp.isinf(_lowerCamelCase      ).any()      )






  def 				UpperCAmelCase__						(							self			:					List[str]      )->Dict:






   '''simple docstring'''

   __lowerCAmelCase					:	int      =       4
   __lowerCAmelCase					:	List[Any]      =       10
   __lowerCAmelCase					:	Union[str, Any]      =       15
   __lowerCAmelCase					:	List[str]      =       2
   __lowerCAmelCase					:	Any      =       1
   __lowerCAmelCase					:	Dict      =       15

   # dummy input_ids and scores
   __lowerCAmelCase					:	Any      =       ids_tensor((batch_size, sequence_length)       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	Optional[int]      =       input_ids.copy()

   __lowerCAmelCase					:	int      =       self._get_uniform_logits(_lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	List[str]      =       scores.copy()

   # instantiate all dist processors
   __lowerCAmelCase					:	Union[str, Any]      =       FlaxTemperatureLogitsWarper(temperature=0.5      )
   __lowerCAmelCase					:	Any      =       FlaxTopKLogitsWarper(3      )
   __lowerCAmelCase					:	Any      =       FlaxTopPLogitsWarper(0.8      )

   # instantiate all logits processors
   __lowerCAmelCase					:	Union[str, Any]      =       FlaxMinLengthLogitsProcessor(min_length=10       ,  eos_token_id=_lowerCamelCase      )
   __lowerCAmelCase					:	List[str]      =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCamelCase      )
   __lowerCAmelCase					:	Optional[Any]      =       FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCamelCase       ,  eos_token_id=_lowerCamelCase      )

   __lowerCAmelCase					:	Optional[int]      =       10

   # no processor list
   __lowerCAmelCase					:	str      =       temp_dist_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   __lowerCAmelCase					:	Union[str, Any]      =       top_k_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   __lowerCAmelCase					:	Union[str, Any]      =       top_p_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   __lowerCAmelCase					:	Dict      =       min_dist_proc(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   __lowerCAmelCase					:	List[str]      =       bos_dist_proc(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
   __lowerCAmelCase					:	Union[str, Any]      =       eos_dist_proc(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )

   # with processor list
   __lowerCAmelCase					:	Optional[int]      =       FlaxLogitsProcessorList(
       [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]      )
   __lowerCAmelCase					:	List[str]      =       processor(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )

   # scores should be equal
   self.assertTrue(jnp.allclose(_lowerCamelCase       ,  _lowerCamelCase       ,  atol=1E-3      )      )

   # input_ids should never be changed
   self.assertListEqual(input_ids.tolist()       ,  input_ids_comp.tolist()      )






  def 				UpperCAmelCase__						(							self			:					int      )->Optional[int]:






   '''simple docstring'''

   __lowerCAmelCase					:	Union[str, Any]      =       4
   __lowerCAmelCase					:	Union[str, Any]      =       10
   __lowerCAmelCase					:	Any      =       15
   __lowerCAmelCase					:	List[Any]      =       2
   __lowerCAmelCase					:	int      =       1
   __lowerCAmelCase					:	Optional[Any]      =       15

   # dummy input_ids and scores
   __lowerCAmelCase					:	Dict      =       ids_tensor((batch_size, sequence_length)       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	Union[str, Any]      =       input_ids.copy()

   __lowerCAmelCase					:	Optional[int]      =       self._get_uniform_logits(_lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	str      =       scores.copy()

   # instantiate all dist processors
   __lowerCAmelCase					:	Tuple      =       FlaxTemperatureLogitsWarper(temperature=0.5      )
   __lowerCAmelCase					:	List[Any]      =       FlaxTopKLogitsWarper(3      )
   __lowerCAmelCase					:	Dict      =       FlaxTopPLogitsWarper(0.8      )

   # instantiate all logits processors
   __lowerCAmelCase					:	int      =       FlaxMinLengthLogitsProcessor(min_length=10       ,  eos_token_id=_lowerCamelCase      )
   __lowerCAmelCase					:	List[str]      =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCamelCase      )
   __lowerCAmelCase					:	Tuple      =       FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCamelCase       ,  eos_token_id=_lowerCamelCase      )

   __lowerCAmelCase					:	str      =       10

   # no processor list
   def run_no_processor_list(_snake_case			:					List[Any]       ,  _snake_case			:					Tuple       ,  _snake_case			:					Optional[int]      ):
    __lowerCAmelCase					:	Any      =       temp_dist_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
    __lowerCAmelCase					:	Optional[int]      =       top_k_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
    __lowerCAmelCase					:	Union[str, Any]      =       top_p_warp(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
    __lowerCAmelCase					:	Optional[int]      =       min_dist_proc(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
    __lowerCAmelCase					:	Dict      =       bos_dist_proc(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
    __lowerCAmelCase					:	Optional[int]      =       eos_dist_proc(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
    return scores

   # with processor list
   def run_processor_list(_snake_case			:					List[Any]       ,  _snake_case			:					Any       ,  _snake_case			:					str      ):
    __lowerCAmelCase					:	str      =       FlaxLogitsProcessorList(
        [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]      )
    __lowerCAmelCase					:	int      =       processor(_lowerCamelCase       ,  _lowerCamelCase       ,  cur_len=_lowerCamelCase      )
    return scores

   __lowerCAmelCase					:	Dict      =       jax.jit(_lowerCamelCase      )
   __lowerCAmelCase					:	Dict      =       jax.jit(_lowerCamelCase      )

   __lowerCAmelCase					:	int      =       jitted_run_no_processor_list(_lowerCamelCase       ,  _lowerCamelCase       ,  _lowerCamelCase      )
   __lowerCAmelCase					:	List[Any]      =       jitted_run_processor_list(_lowerCamelCase       ,  _lowerCamelCase       ,  _lowerCamelCase      )

   # scores should be equal
   self.assertTrue(jnp.allclose(_lowerCamelCase       ,  _lowerCamelCase       ,  atol=1E-3      )      )

   # input_ids should never be changed
   self.assertListEqual(input_ids.tolist()       ,  input_ids_comp.tolist()      )








from datetime import datetime

import requests






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:str   )      ->   bytes:
 __lowerCAmelCase					:	List[Any]      =       """https://downloadgram.net/wp-json/wppress/video-downloader/video?url="""
 __lowerCAmelCase					:	Dict      =       requests.get(base_url + url   ).json()[0]["""urls"""][0]["""src"""]
 return requests.get(SCREAMING_SNAKE_CASE   ).content


if __name__ == "__main__":
      _UpperCAmelCase             =							input('Enter Video/IGTV url: ').strip()
      _UpperCAmelCase             =							f'''{datetime.now():%Y-%m-%d_%H:%M:%S}.mp4'''
      with open(file_name, 'wb') as fp:
            fp.write(download_video(url))
      print(f'''Done. Video saved to disk as {file_name}.''')