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






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





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






					def 	_snake_case    (   self		):


											"""simple docstring"""





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

											lowercase_	:     Optional[int]													=       self._get_uniform_logits(batch_size=2				,							length=__SCREAMING_SNAKE_CASE		)

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

											# compute softmax
											lowercase_	:     Any													=       jax.nn.softmax(__SCREAMING_SNAKE_CASE				,							axis=-1		)

											lowercase_	:     Tuple													=       FlaxTemperatureLogitsWarper(temperature=0.5		)
											lowercase_	:     Optional[Any]													=       FlaxTemperatureLogitsWarper(temperature=1.3		)

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


											"""simple docstring"""





											lowercase_	:     Any													=       None
											lowercase_	:     List[str]													=       10
											lowercase_	:     Tuple													=       2

											# create ramp distribution
											lowercase_	:     int													=       np.broadcast_to(np.arange(__SCREAMING_SNAKE_CASE		)[None, :]				,							(batch_size, vocab_size)		).copy()
											lowercase_	:     Optional[Any]													=       ramp_logits[1:, : vocab_size // 2] + vocab_size

											lowercase_	:     Optional[int]													=       FlaxTopKLogitsWarper(3		)

											lowercase_	:     Tuple													=       top_k_warp(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)

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

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

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






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Optional[int]													=       None
											lowercase_	:     Optional[int]													=       10
											lowercase_	:     Any													=       2

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

											lowercase_	:     List[Any]													=       FlaxTopPLogitsWarper(0.8		)
											lowercase_	:     Optional[int]													=       np.exp(top_p_warp(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)		)

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

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

											# make ramp_logits more extreme
											lowercase_	:     Union[str, Any]													=       ramp_logits[1] * 100.0

											# make sure at least 2 tokens are kept
											lowercase_	:     Tuple													=       FlaxTopPLogitsWarper(0.9				,							min_tokens_to_keep=2				,							filter_value=0.0		)
											lowercase_	:     Optional[Any]													=       top_p_warp(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)

											# 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 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     str													=       20
											lowercase_	:     Tuple													=       4
											lowercase_	:     Optional[Any]													=       0

											lowercase_	:     Optional[Any]													=       FlaxMinLengthLogitsProcessor(min_length=10				,							eos_token_id=__SCREAMING_SNAKE_CASE		)

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

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






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Any													=       20
											lowercase_	:     Tuple													=       4
											lowercase_	:     str													=       0

											lowercase_	:     Tuple													=       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE		)

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






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Union[str, Any]													=       20
											lowercase_	:     int													=       4
											lowercase_	:     Tuple													=       0
											lowercase_	:     Any													=       5

											lowercase_	:     List[Any]													=       FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE				,							eos_token_id=__SCREAMING_SNAKE_CASE		)

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






					def 	_snake_case    (   self		):


											"""simple docstring"""





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

											# dummy input_ids and scores
											lowercase_	:     Dict													=       ids_tensor((batch_size, sequence_length)				,							__SCREAMING_SNAKE_CASE		)
											lowercase_	:     str													=       input_ids.copy()

											lowercase_	:     Dict													=       self._get_uniform_logits(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[int]													=       scores.copy()

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

											# instantiate all logits processors
											lowercase_	:     Union[str, Any]													=       FlaxMinLengthLogitsProcessor(min_length=10				,							eos_token_id=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[int]													=       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Tuple													=       FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE				,							eos_token_id=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Tuple													=       10

											# no processor list
											lowercase_	:     Optional[Any]													=       temp_dist_warp(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     int													=       top_k_warp(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Tuple													=       top_p_warp(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Any													=       min_dist_proc(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Dict													=       bos_dist_proc(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Any													=       eos_dist_proc(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)

											# with processor list
											lowercase_	:     Any													=       FlaxLogitsProcessorList(
											    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
											lowercase_	:     Tuple													=       processor(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)

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

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






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Optional[int]													=       4
											lowercase_	:     Optional[Any]													=       10
											lowercase_	:     List[str]													=       15
											lowercase_	:     int													=       2
											lowercase_	:     Optional[int]													=       1
											lowercase_	:     int													=       15

											# dummy input_ids and scores
											lowercase_	:     Optional[int]													=       ids_tensor((batch_size, sequence_length)				,							__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Dict													=       input_ids.copy()

											lowercase_	:     Optional[Any]													=       self._get_uniform_logits(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Union[str, Any]													=       scores.copy()

											# instantiate all dist processors
											lowercase_	:     List[str]													=       FlaxTemperatureLogitsWarper(temperature=0.5		)
											lowercase_	:     Optional[int]													=       FlaxTopKLogitsWarper(3		)
											lowercase_	:     int													=       FlaxTopPLogitsWarper(0.8		)

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

											lowercase_	:     Optional[int]													=       10

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

											# with processor list
											def run_processor_list(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																	lowercase_	:     Optional[int]													=       FlaxLogitsProcessorList(
																	    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
																	lowercase_	:     Tuple													=       processor(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							cur_len=__SCREAMING_SNAKE_CASE		)
																	return scores

											lowercase_	:     List[Any]													=       jax.jit(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Any													=       jax.jit(__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[Any]													=       jitted_run_no_processor_list(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)
											lowercase_	:     int													=       jitted_run_processor_list(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)

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

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





'''simple docstring'''







import unittest

from knapsack import greedy_knapsack as kp






class  lowerCAmelCase__   (  unittest.TestCase     ):






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     List[str]													=       [10, 20, 30, 40, 50, 60]
											lowercase_	:     Optional[Any]													=       [2, 4, 6, 8, 10, 12]
											lowercase_	:     Union[str, Any]													=       1_00
											self.assertEqual(kp.calc_profit(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)				,							2_10		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											self.assertRaisesRegex(__SCREAMING_SNAKE_CASE				,							'''max_weight must greater than zero.'''		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											self.assertRaisesRegex(__SCREAMING_SNAKE_CASE				,							'''Weight can not be negative.'''		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											self.assertRaisesRegex(__SCREAMING_SNAKE_CASE				,							'''Profit can not be negative.'''		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											self.assertRaisesRegex(__SCREAMING_SNAKE_CASE				,							'''max_weight must greater than zero.'''		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											self.assertRaisesRegex(
											    __SCREAMING_SNAKE_CASE				,							'''The length of profit and weight must be same.'''		)


if __name__ == "__main__":
							unittest.main()



