

'''simple docstring'''



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       SCREAMING_SNAKE_CASE       (				unittest.TestCase       ):






      '''simple docstring'''






      def 						__UpperCAmelCase      (       self			:		Tuple  ,       snake_case			:		int  ,       snake_case			:		int ):



          """simple docstring"""





          _snake_case			:			List[Any]				       =		jnp.ones((batch_size, length) ) / length
          return scores




      def 						__UpperCAmelCase      (       self			:		Dict ):



          """simple docstring"""





          _snake_case			:			List[str]				       =		None
          _snake_case			:			str				       =		20

          _snake_case			:			int				       =		self._get_uniform_logits(batch_size=2  ,       length=UpperCAmelCase__ )

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

          # compute softmax
          _snake_case			:			str				       =		jax.nn.softmax(UpperCAmelCase__  ,       axis=-1 )

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

          _snake_case			:			Optional[int]				       =		jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase__  ,       scores.copy()  ,       cur_len=UpperCAmelCase__ )  ,       axis=-1 )
          _snake_case			:			Optional[Any]				       =		jax.nn.softmax(temp_dist_warper_smoother(UpperCAmelCase__  ,       scores.copy()  ,       cur_len=UpperCAmelCase__ )  ,       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			:		Optional[Any] ):



          """simple docstring"""





          _snake_case			:			Union[str, Any]				       =		None
          _snake_case			:			Optional[Any]				       =		10
          _snake_case			:			Any				       =		2

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

          _snake_case			:			str				       =		FlaxTopKLogitsWarper(3 )

          _snake_case			:			List[Any]				       =		top_k_warp(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )

          # 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
          _snake_case			:			Optional[int]				       =		5
          _snake_case			:			Tuple				       =		FlaxTopKLogitsWarper(top_k=1  ,       filter_value=0.0  ,       min_tokens_to_keep=3 )

          _snake_case			:			Optional[int]				       =		np.broadcast_to(np.arange(UpperCAmelCase__ )[None, :]  ,       (batch_size, length) ).copy()
          _snake_case			:			int				       =		top_k_warp_safety_check(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )

          # 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			:		Dict ):



          """simple docstring"""





          _snake_case			:			Any				       =		None
          _snake_case			:			Optional[Any]				       =		10
          _snake_case			:			Dict				       =		2

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

          _snake_case			:			Tuple				       =		FlaxTopPLogitsWarper(0.8 )
          _snake_case			:			int				       =		np.exp(top_p_warp(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ ) )

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

          # check edge cases with negative and extreme logits
          _snake_case			:			Tuple				       =		np.broadcast_to(np.arange(UpperCAmelCase__ )[None, :]  ,       (batch_size, vocab_size) ).copy() - (
              vocab_size // 2
          )

          # make ramp_logits more extreme
          _snake_case			:			List[Any]				       =		ramp_logits[1] * 100.0

          # make sure at least 2 tokens are kept
          _snake_case			:			Union[str, Any]				       =		FlaxTopPLogitsWarper(0.9  ,       min_tokens_to_keep=2  ,       filter_value=0.0 )
          _snake_case			:			List[str]				       =		top_p_warp(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )

          # 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] ):



          """simple docstring"""





          _snake_case			:			Dict				       =		20
          _snake_case			:			Any				       =		4
          _snake_case			:			Tuple				       =		0

          _snake_case			:			List[Any]				       =		FlaxMinLengthLogitsProcessor(min_length=10  ,       eos_token_id=UpperCAmelCase__ )

          # check that min length is applied at length 5
          _snake_case			:			List[str]				       =		ids_tensor((batch_size, 20)  ,       vocab_size=20 )
          _snake_case			:			List[Any]				       =		5
          _snake_case			:			Optional[Any]				       =		self._get_uniform_logits(UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			int				       =		min_dist_processor(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()  ,       4 * [-float('inf' )] )

          # check that min length is not applied anymore at length 15
          _snake_case			:			Tuple				       =		self._get_uniform_logits(UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			Optional[int]				       =		15
          _snake_case			:			List[Any]				       =		min_dist_processor(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          self.assertFalse(jnp.isinf(UpperCAmelCase__ ).any() )




      def 						__UpperCAmelCase      (       self			:		Optional[int] ):



          """simple docstring"""





          _snake_case			:			Tuple				       =		20
          _snake_case			:			Any				       =		4
          _snake_case			:			Union[str, Any]				       =		0

          _snake_case			:			Optional[int]				       =		FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__ )

          # check that all scores are -inf except the bos_token_id score
          _snake_case			:			int				       =		ids_tensor((batch_size, 1)  ,       vocab_size=20 )
          _snake_case			:			str				       =		1
          _snake_case			:			Any				       =		self._get_uniform_logits(UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			str				       =		logits_processor(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          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
          _snake_case			:			Optional[Any]				       =		3
          _snake_case			:			List[Any]				       =		self._get_uniform_logits(UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			Union[str, Any]				       =		logits_processor(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          self.assertFalse(jnp.isinf(UpperCAmelCase__ ).any() )




      def 						__UpperCAmelCase      (       self			:		int ):



          """simple docstring"""





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

          _snake_case			:			Dict				       =		FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__  ,       eos_token_id=UpperCAmelCase__ )

          # check that all scores are -inf except the eos_token_id when max_length is reached
          _snake_case			:			int				       =		ids_tensor((batch_size, 4)  ,       vocab_size=20 )
          _snake_case			:			Any				       =		4
          _snake_case			:			List[str]				       =		self._get_uniform_logits(UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			Optional[int]				       =		logits_processor(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          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
          _snake_case			:			List[str]				       =		3
          _snake_case			:			Union[str, Any]				       =		self._get_uniform_logits(UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			Any				       =		logits_processor(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          self.assertFalse(jnp.isinf(UpperCAmelCase__ ).any() )




      def 						__UpperCAmelCase      (       self			:		Dict ):



          """simple docstring"""





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

          # dummy input_ids and scores
          _snake_case			:			List[str]				       =		ids_tensor((batch_size, sequence_length)  ,       UpperCAmelCase__ )
          _snake_case			:			Any				       =		input_ids.copy()

          _snake_case			:			Optional[Any]				       =		self._get_uniform_logits(UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			List[Any]				       =		scores.copy()

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

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

          _snake_case			:			Dict				       =		10

          # no processor list
          _snake_case			:			Optional[int]				       =		temp_dist_warp(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          _snake_case			:			Tuple				       =		top_k_warp(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          _snake_case			:			Any				       =		top_p_warp(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          _snake_case			:			List[Any]				       =		min_dist_proc(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          _snake_case			:			List[Any]				       =		bos_dist_proc(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
          _snake_case			:			Any				       =		eos_dist_proc(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )

          # with processor list
          _snake_case			:			Optional[int]				       =		FlaxLogitsProcessorList(
              [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
          _snake_case			:			Tuple				       =		processor(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )

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

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




      def 						__UpperCAmelCase      (       self			:		Dict ):



          """simple docstring"""





          _snake_case			:			Tuple				       =		4
          _snake_case			:			str				       =		10
          _snake_case			:			List[str]				       =		15
          _snake_case			:			List[Any]				       =		2
          _snake_case			:			List[Any]				       =		1
          _snake_case			:			List[Any]				       =		15

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

          _snake_case			:			List[str]				       =		self._get_uniform_logits(UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			List[Any]				       =		scores.copy()

          # instantiate all dist processors
          _snake_case			:			Dict				       =		FlaxTemperatureLogitsWarper(temperature=0.5 )
          _snake_case			:			Dict				       =		FlaxTopKLogitsWarper(3 )
          _snake_case			:			int				       =		FlaxTopPLogitsWarper(0.8 )

          # instantiate all logits processors
          _snake_case			:			Any				       =		FlaxMinLengthLogitsProcessor(min_length=10  ,       eos_token_id=UpperCAmelCase__ )
          _snake_case			:			Dict				       =		FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__ )
          _snake_case			:			Dict				       =		FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__  ,       eos_token_id=UpperCAmelCase__ )

          _snake_case			:			List[str]				       =		10

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

          # with processor list
          def run_processor_list(snake_case			:		Union[str, Any]  ,       snake_case			:		str  ,       snake_case			:		Dict ):
              _snake_case			:			int				       =		FlaxLogitsProcessorList(
                  [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
              _snake_case			:			Dict				       =		processor(UpperCAmelCase__  ,       UpperCAmelCase__  ,       cur_len=UpperCAmelCase__ )
              return scores

          _snake_case			:			List[Any]				       =		jax.jit(UpperCAmelCase__ )
          _snake_case			:			List[Any]				       =		jax.jit(UpperCAmelCase__ )

          _snake_case			:			List[Any]				       =		jitted_run_no_processor_list(UpperCAmelCase__  ,       UpperCAmelCase__  ,       UpperCAmelCase__ )
          _snake_case			:			Dict				       =		jitted_run_processor_list(UpperCAmelCase__  ,       UpperCAmelCase__  ,       UpperCAmelCase__ )

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

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









from __future__ import annotations

import requests







def SCREAMING_SNAKE_CASE							(			lowercase_		)  ->			dict:



       """simple docstring"""

       A__ 					=       f"""https://hacker-news.firebaseio.com/v0/item/{story_id}.json?print=pretty"""
       return requests.get(lowercase_		).json()







def SCREAMING_SNAKE_CASE							(			lowercase_ = 10		)  ->			list[dict]:



       """simple docstring"""

       A__ 					=       '''https://hacker-news.firebaseio.com/v0/topstories.json?print=pretty'''
       A__ 					=       requests.get(lowercase_		).json()[:max_stories]
       return [get_hackernews_story(lowercase_		) for story_id in story_ids]







def SCREAMING_SNAKE_CASE							(			lowercase_ = 10		)  ->			str:



       """simple docstring"""

       A__ 					=       hackernews_top_stories(lowercase_		)
       return "\n".join('''* [{title}]({url})'''.format(**lowercase_		) for story in stories		)


if __name__ == "__main__":
  print(hackernews_top_stories_as_markdown())


