



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







       '''simple docstring'''


       def 	lowercase__							(				self     :    Optional[int]						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    int       )					->		int:


             '''simple docstring'''




             UpperCAmelCase_     					=				jnp.ones((batch_size, length)       ) / length
             return scores



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


             '''simple docstring'''




             UpperCAmelCase_     					=				None
             UpperCAmelCase_     					=				20

             UpperCAmelCase_     					=				self._get_uniform_logits(batch_size=2						,							length=_UpperCAmelCase       )

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

             # compute softmax
             UpperCAmelCase_     					=				jax.nn.softmax(_UpperCAmelCase						,							axis=-1       )

             UpperCAmelCase_     					=				FlaxTemperatureLogitsWarper(temperature=0.5       )
             UpperCAmelCase_     					=				FlaxTemperatureLogitsWarper(temperature=1.3       )

             UpperCAmelCase_     					=				jax.nn.softmax(temp_dist_warper_sharper(_UpperCAmelCase						,							scores.copy()						,							cur_len=_UpperCAmelCase       )						,							axis=-1       )
             UpperCAmelCase_     					=				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 	lowercase__							(				self     :    Dict       )					->		Any:


             '''simple docstring'''




             UpperCAmelCase_     					=				None
             UpperCAmelCase_     					=				10
             UpperCAmelCase_     					=				2

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

             UpperCAmelCase_     					=				FlaxTopKLogitsWarper(3       )

             UpperCAmelCase_     					=				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
             UpperCAmelCase_     					=				5
             UpperCAmelCase_     					=				FlaxTopKLogitsWarper(top_k=1						,							filter_value=0.0						,							min_tokens_to_keep=3       )

             UpperCAmelCase_     					=				np.broadcast_to(np.arange(_UpperCAmelCase       )[None, :]						,							(batch_size, length)       ).copy()
             UpperCAmelCase_     					=				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 	lowercase__							(				self     :    Dict       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				None
             UpperCAmelCase_     					=				10
             UpperCAmelCase_     					=				2

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

             UpperCAmelCase_     					=				FlaxTopPLogitsWarper(0.8       )
             UpperCAmelCase_     					=				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
             UpperCAmelCase_     					=				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
             UpperCAmelCase_     					=				np.broadcast_to(np.arange(_UpperCAmelCase       )[None, :]						,							(batch_size, vocab_size)       ).copy() - (
                 vocab_size // 2
             )

             # make ramp_logits more extreme
             UpperCAmelCase_     					=				ramp_logits[1] * 100.0

             # make sure at least 2 tokens are kept
             UpperCAmelCase_     					=				FlaxTopPLogitsWarper(0.9						,							min_tokens_to_keep=2						,							filter_value=0.0       )
             UpperCAmelCase_     					=				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 	lowercase__							(				self     :    Optional[int]       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				20
             UpperCAmelCase_     					=				4
             UpperCAmelCase_     					=				0

             UpperCAmelCase_     					=				FlaxMinLengthLogitsProcessor(min_length=10						,							eos_token_id=_UpperCAmelCase       )

             # check that min length is applied at length 5
             UpperCAmelCase_     					=				ids_tensor((batch_size, 20)						,							vocab_size=20       )
             UpperCAmelCase_     					=				5
             UpperCAmelCase_     					=				self._get_uniform_logits(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				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
             UpperCAmelCase_     					=				self._get_uniform_logits(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				15
             UpperCAmelCase_     					=				min_dist_processor(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
             self.assertFalse(jnp.isinf(_UpperCAmelCase       ).any()       )



       def 	lowercase__							(				self     :    Union[str, Any]       )					->		str:


             '''simple docstring'''




             UpperCAmelCase_     					=				20
             UpperCAmelCase_     					=				4
             UpperCAmelCase_     					=				0

             UpperCAmelCase_     					=				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_UpperCAmelCase       )

             # check that all scores are -inf except the bos_token_id score
             UpperCAmelCase_     					=				ids_tensor((batch_size, 1)						,							vocab_size=20       )
             UpperCAmelCase_     					=				1
             UpperCAmelCase_     					=				self._get_uniform_logits(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				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
             UpperCAmelCase_     					=				3
             UpperCAmelCase_     					=				self._get_uniform_logits(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				logits_processor(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
             self.assertFalse(jnp.isinf(_UpperCAmelCase       ).any()       )



       def 	lowercase__							(				self     :    Any       )					->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				20
             UpperCAmelCase_     					=				4
             UpperCAmelCase_     					=				0
             UpperCAmelCase_     					=				5

             UpperCAmelCase_     					=				FlaxForcedEOSTokenLogitsProcessor(max_length=_UpperCAmelCase						,							eos_token_id=_UpperCAmelCase       )

             # check that all scores are -inf except the eos_token_id when max_length is reached
             UpperCAmelCase_     					=				ids_tensor((batch_size, 4)						,							vocab_size=20       )
             UpperCAmelCase_     					=				4
             UpperCAmelCase_     					=				self._get_uniform_logits(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				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
             UpperCAmelCase_     					=				3
             UpperCAmelCase_     					=				self._get_uniform_logits(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				logits_processor(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
             self.assertFalse(jnp.isinf(_UpperCAmelCase       ).any()       )



       def 	lowercase__							(				self     :    Union[str, Any]       )					->		int:


             '''simple docstring'''




             UpperCAmelCase_     					=				4
             UpperCAmelCase_     					=				10
             UpperCAmelCase_     					=				15
             UpperCAmelCase_     					=				2
             UpperCAmelCase_     					=				1
             UpperCAmelCase_     					=				15

             # dummy input_ids and scores
             UpperCAmelCase_     					=				ids_tensor((batch_size, sequence_length)						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				input_ids.copy()

             UpperCAmelCase_     					=				self._get_uniform_logits(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				scores.copy()

             # instantiate all dist processors
             UpperCAmelCase_     					=				FlaxTemperatureLogitsWarper(temperature=0.5       )
             UpperCAmelCase_     					=				FlaxTopKLogitsWarper(3       )
             UpperCAmelCase_     					=				FlaxTopPLogitsWarper(0.8       )

             # instantiate all logits processors
             UpperCAmelCase_     					=				FlaxMinLengthLogitsProcessor(min_length=10						,							eos_token_id=_UpperCAmelCase       )
             UpperCAmelCase_     					=				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_UpperCAmelCase       )
             UpperCAmelCase_     					=				FlaxForcedEOSTokenLogitsProcessor(max_length=_UpperCAmelCase						,							eos_token_id=_UpperCAmelCase       )

             UpperCAmelCase_     					=				10

             # no processor list
             UpperCAmelCase_     					=				temp_dist_warp(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
             UpperCAmelCase_     					=				top_k_warp(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
             UpperCAmelCase_     					=				top_p_warp(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
             UpperCAmelCase_     					=				min_dist_proc(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
             UpperCAmelCase_     					=				bos_dist_proc(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
             UpperCAmelCase_     					=				eos_dist_proc(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )

             # with processor list
             UpperCAmelCase_     					=				FlaxLogitsProcessorList(
                 [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]       )
             UpperCAmelCase_     					=				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 	lowercase__							(				self     :    List[Any]       )					->		str:


             '''simple docstring'''




             UpperCAmelCase_     					=				4
             UpperCAmelCase_     					=				10
             UpperCAmelCase_     					=				15
             UpperCAmelCase_     					=				2
             UpperCAmelCase_     					=				1
             UpperCAmelCase_     					=				15

             # dummy input_ids and scores
             UpperCAmelCase_     					=				ids_tensor((batch_size, sequence_length)						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				input_ids.copy()

             UpperCAmelCase_     					=				self._get_uniform_logits(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				scores.copy()

             # instantiate all dist processors
             UpperCAmelCase_     					=				FlaxTemperatureLogitsWarper(temperature=0.5       )
             UpperCAmelCase_     					=				FlaxTopKLogitsWarper(3       )
             UpperCAmelCase_     					=				FlaxTopPLogitsWarper(0.8       )

             # instantiate all logits processors
             UpperCAmelCase_     					=				FlaxMinLengthLogitsProcessor(min_length=10						,							eos_token_id=_UpperCAmelCase       )
             UpperCAmelCase_     					=				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_UpperCAmelCase       )
             UpperCAmelCase_     					=				FlaxForcedEOSTokenLogitsProcessor(max_length=_UpperCAmelCase						,							eos_token_id=_UpperCAmelCase       )

             UpperCAmelCase_     					=				10

             # no processor list
             def run_no_processor_list(_UpperCAmelCase     :    Optional[int]						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    str       ):
                   UpperCAmelCase_     					=				temp_dist_warp(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
                   UpperCAmelCase_     					=				top_k_warp(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
                   UpperCAmelCase_     					=				top_p_warp(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
                   UpperCAmelCase_     					=				min_dist_proc(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
                   UpperCAmelCase_     					=				bos_dist_proc(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
                   UpperCAmelCase_     					=				eos_dist_proc(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
                   return scores

             # with processor list
             def run_processor_list(_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Any       ):
                   UpperCAmelCase_     					=				FlaxLogitsProcessorList(
                       [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]       )
                   UpperCAmelCase_     					=				processor(_UpperCAmelCase						,							_UpperCAmelCase						,							cur_len=_UpperCAmelCase       )
                   return scores

             UpperCAmelCase_     					=				jax.jit(_UpperCAmelCase       )
             UpperCAmelCase_     					=				jax.jit(_UpperCAmelCase       )

             UpperCAmelCase_     					=				jitted_run_no_processor_list(_UpperCAmelCase						,							_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				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()       )






'''simple docstring'''
import requests

lowercase__											=''  # <-- Put your OpenWeatherMap appid here!
lowercase__											='https://api.openweathermap.org/data/2.5/'






def      UpperCamelCase_				(      A__ = "Chicago"  ,      A__ = APPID						):
	return requests.get(URL_BASE + """weather"""  ,      params=locals()						).json()






def      UpperCamelCase_				(      A__ = "Kolkata, India"  ,      A__ = APPID						):
	return requests.get(URL_BASE + """forecast"""  ,      params=locals()						).json()






def      UpperCamelCase_				(      A__ = 55.68  ,      A__ = 12.57  ,      A__ = APPID						):
	return requests.get(URL_BASE + """onecall"""  ,      params=locals()						).json()


if __name__ == "__main__":
		from pprint import pprint

		while True:
				lowercase__											=input('Enter a location:').strip()
				if location:
						pprint(current_weather(location))
				else:
						break
