
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 		A__					(				self	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		int			=						jnp.ones((batch_size, length)) / length
			return scores


		def 		A__					(				self):
			_UpperCamelCase   :		Optional[Any]			=						None
			_UpperCamelCase   :		Tuple			=						20

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

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

			# compute softmax
			_UpperCamelCase   :		str			=						jax.nn.softmax(__snake_case	,       axis=-1)

			_UpperCamelCase   :		Tuple			=						FlaxTemperatureLogitsWarper(temperature=0.5)
			_UpperCamelCase   :		Optional[int]			=						FlaxTemperatureLogitsWarper(temperature=1.3)

			_UpperCamelCase   :		Optional[int]			=						jax.nn.softmax(temp_dist_warper_sharper(__snake_case	,       scores.copy()	,       cur_len=__snake_case)	,       axis=-1)
			_UpperCamelCase   :		Any			=						jax.nn.softmax(temp_dist_warper_smoother(__snake_case	,       scores.copy()	,       cur_len=__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 		A__					(				self):
			_UpperCamelCase   :		Optional[int]			=						None
			_UpperCamelCase   :		List[str]			=						10
			_UpperCamelCase   :		Optional[int]			=						2

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

			_UpperCamelCase   :		List[Any]			=						FlaxTopKLogitsWarper(3)

			_UpperCamelCase   :		List[Any]			=						top_k_warp(__snake_case	,       __snake_case	,       cur_len=__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
			_UpperCamelCase   :		Optional[int]			=						5
			_UpperCamelCase   :		str			=						FlaxTopKLogitsWarper(top_k=1	,       filter_value=0.0	,       min_tokens_to_keep=3)

			_UpperCamelCase   :		str			=						np.broadcast_to(np.arange(__snake_case)[None, :]	,       (batch_size, length)).copy()
			_UpperCamelCase   :		str			=						top_k_warp_safety_check(__snake_case	,       __snake_case	,       cur_len=__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 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						None
			_UpperCamelCase   :		Any			=						10
			_UpperCamelCase   :		str			=						2

			# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
			_UpperCamelCase   :		Union[str, Any]			=						np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.1_5, 0.3, 0.3, 0.2_5]]))

			_UpperCamelCase   :		Optional[Any]			=						FlaxTopPLogitsWarper(0.8)
			_UpperCamelCase   :		List[str]			=						np.exp(top_p_warp(__snake_case	,       __snake_case	,       cur_len=__snake_case))

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

			# check edge cases with negative and extreme logits
			_UpperCamelCase   :		List[Any]			=						np.broadcast_to(np.arange(__snake_case)[None, :]	,       (batch_size, vocab_size)).copy() - (
			    vocab_size // 2
			)

			# make ramp_logits more extreme
			_UpperCamelCase   :		List[Any]			=						ramp_logits[1] * 1_0_0.0

			# make sure at least 2 tokens are kept
			_UpperCamelCase   :		Optional[int]			=						FlaxTopPLogitsWarper(0.9	,       min_tokens_to_keep=2	,       filter_value=0.0)
			_UpperCamelCase   :		List[str]			=						top_p_warp(__snake_case	,       __snake_case	,       cur_len=__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 		A__					(				self):
			_UpperCamelCase   :		Optional[Any]			=						20
			_UpperCamelCase   :		Any			=						4
			_UpperCamelCase   :		Optional[int]			=						0

			_UpperCamelCase   :		Tuple			=						FlaxMinLengthLogitsProcessor(min_length=10	,       eos_token_id=__snake_case)

			# check that min length is applied at length 5
			_UpperCamelCase   :		Dict			=						ids_tensor((batch_size, 20)	,       vocab_size=20)
			_UpperCamelCase   :		str			=						5
			_UpperCamelCase   :		int			=						self._get_uniform_logits(__snake_case	,       __snake_case)
			_UpperCamelCase   :		Dict			=						min_dist_processor(__snake_case	,       __snake_case	,       cur_len=__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
			_UpperCamelCase   :		Union[str, Any]			=						self._get_uniform_logits(__snake_case	,       __snake_case)
			_UpperCamelCase   :		int			=						15
			_UpperCamelCase   :		Optional[int]			=						min_dist_processor(__snake_case	,       __snake_case	,       cur_len=__snake_case)
			self.assertFalse(jnp.isinf(__snake_case).any())


		def 		A__					(				self):
			_UpperCamelCase   :		int			=						20
			_UpperCamelCase   :		List[Any]			=						4
			_UpperCamelCase   :		Optional[int]			=						0

			_UpperCamelCase   :		Dict			=						FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__snake_case)

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


		def 		A__					(				self):
			_UpperCamelCase   :		Any			=						20
			_UpperCamelCase   :		Any			=						4
			_UpperCamelCase   :		str			=						0
			_UpperCamelCase   :		List[Any]			=						5

			_UpperCamelCase   :		Any			=						FlaxForcedEOSTokenLogitsProcessor(max_length=__snake_case	,       eos_token_id=__snake_case)

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


		def 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						4
			_UpperCamelCase   :		int			=						10
			_UpperCamelCase   :		int			=						15
			_UpperCamelCase   :		Tuple			=						2
			_UpperCamelCase   :		Optional[int]			=						1
			_UpperCamelCase   :		Dict			=						15

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

			_UpperCamelCase   :		Union[str, Any]			=						self._get_uniform_logits(__snake_case	,       __snake_case)
			_UpperCamelCase   :		int			=						scores.copy()

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

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

			_UpperCamelCase   :		List[Any]			=						10

			# no processor list
			_UpperCamelCase   :		str			=						temp_dist_warp(__snake_case	,       __snake_case	,       cur_len=__snake_case)
			_UpperCamelCase   :		Tuple			=						top_k_warp(__snake_case	,       __snake_case	,       cur_len=__snake_case)
			_UpperCamelCase   :		Tuple			=						top_p_warp(__snake_case	,       __snake_case	,       cur_len=__snake_case)
			_UpperCamelCase   :		Dict			=						min_dist_proc(__snake_case	,       __snake_case	,       cur_len=__snake_case)
			_UpperCamelCase   :		Dict			=						bos_dist_proc(__snake_case	,       __snake_case	,       cur_len=__snake_case)
			_UpperCamelCase   :		int			=						eos_dist_proc(__snake_case	,       __snake_case	,       cur_len=__snake_case)

			# with processor list
			_UpperCamelCase   :		List[str]			=						FlaxLogitsProcessorList(
			    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
			_UpperCamelCase   :		Optional[Any]			=						processor(__snake_case	,       __snake_case	,       cur_len=__snake_case)

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

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






		def 		A__					(				self):
			_UpperCamelCase   :		int			=						4
			_UpperCamelCase   :		Dict			=						10
			_UpperCamelCase   :		Optional[Any]			=						15
			_UpperCamelCase   :		Union[str, Any]			=						2
			_UpperCamelCase   :		Dict			=						1
			_UpperCamelCase   :		Tuple			=						15

			# dummy input_ids and scores
			_UpperCamelCase   :		Any			=						ids_tensor((batch_size, sequence_length)	,       __snake_case)
			_UpperCamelCase   :		Any			=						input_ids.copy()

			_UpperCamelCase   :		Dict			=						self._get_uniform_logits(__snake_case	,       __snake_case)
			_UpperCamelCase   :		Tuple			=						scores.copy()

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

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

			_UpperCamelCase   :		List[Any]			=						10

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

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

			_UpperCamelCase   :		Dict			=						jax.jit(__snake_case)
			_UpperCamelCase   :		Union[str, Any]			=						jax.jit(__snake_case)

			_UpperCamelCase   :		str			=						jitted_run_no_processor_list(__snake_case	,       __snake_case	,       __snake_case)
			_UpperCamelCase   :		Any			=						jitted_run_processor_list(__snake_case	,       __snake_case	,       __snake_case)

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

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







from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase__								=		logging.get_logger(__name__)

lowerCAmelCase__								=		{
    """RWKV/rwkv-4-169m-pile""": """https://huggingface.co/RWKV/rwkv-4-169m-pile/resolve/main/config.json""",
    """RWKV/rwkv-4-430m-pile""": """https://huggingface.co/RWKV/rwkv-4-430m-pile/resolve/main/config.json""",
    """RWKV/rwkv-4-1b5-pile""": """https://huggingface.co/RWKV/rwkv-4-1b5-pile/resolve/main/config.json""",
    """RWKV/rwkv-4-3b-pile""": """https://huggingface.co/RWKV/rwkv-4-3b-pile/resolve/main/config.json""",
    """RWKV/rwkv-4-7b-pile""": """https://huggingface.co/RWKV/rwkv-4-7b-pile/resolve/main/config.json""",
    """RWKV/rwkv-4-14b-pile""": """https://huggingface.co/RWKV/rwkv-4-14b-pile/resolve/main/config.json""",
    """RWKV/rwkv-raven-1b5""": """https://huggingface.co/RWKV/rwkv-raven-1b5/resolve/main/config.json""",
    """RWKV/rwkv-raven-3b""": """https://huggingface.co/RWKV/rwkv-raven-3b/resolve/main/config.json""",
    """RWKV/rwkv-raven-7b""": """https://huggingface.co/RWKV/rwkv-raven-7b/resolve/main/config.json""",
    """RWKV/rwkv-raven-14b""": """https://huggingface.co/RWKV/rwkv-raven-14b/resolve/main/config.json""",
}







class 	lowercase    (				_lowercase							):


		"""simple docstring"""




		a__					       =  "rwkv"
		a__					       =  {"max_position_embeddings": "context_length"}


		def __init__(				self	,       __snake_case=5_02_77	,       __snake_case=10_24	,       __snake_case=40_96	,       __snake_case=32	,       __snake_case=None	,       __snake_case=None	,       __snake_case=1e-5	,       __snake_case=0	,       __snake_case=0	,       __snake_case=6	,       __snake_case=False	,       __snake_case=True	,       **__snake_case	,       ):
			_UpperCamelCase   :		str			=						vocab_size
			_UpperCamelCase   :		int			=						context_length
			_UpperCamelCase   :		Tuple			=						hidden_size
			_UpperCamelCase   :		Tuple			=						num_hidden_layers
			_UpperCamelCase   :		Dict			=						attention_hidden_size if attention_hidden_size is not None else hidden_size
			_UpperCamelCase   :		Tuple			=						intermediate_size if intermediate_size is not None else 4 * hidden_size
			_UpperCamelCase   :		Union[str, Any]			=						layer_norm_epsilon
			_UpperCamelCase   :		Dict			=						rescale_every
			_UpperCamelCase   :		Optional[Any]			=						use_cache

			_UpperCamelCase   :		str			=						bos_token_id
			_UpperCamelCase   :		Optional[Any]			=						eos_token_id

			super().__init__(
			    tie_word_embeddings=__snake_case	,       bos_token_id=__snake_case	,       eos_token_id=__snake_case	,       **__snake_case)




