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






     def    SCREAMING_SNAKE_CASE__   ( self		:		Any  ,				a		:		int  ,				a		:		int    ):





         """simple docstring"""



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






     def    SCREAMING_SNAKE_CASE__   ( self		:		Optional[Any]    ):





         """simple docstring"""



         __lowerCamelCase		=  None
         __lowerCamelCase		=  20

         __lowerCamelCase		=  self._get_uniform_logits(batch_size=2  ,				length=a    )

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

         # compute softmax
         __lowerCamelCase		=  jax.nn.softmax(a  ,				axis=-1    )

         __lowerCamelCase		=  FlaxTemperatureLogitsWarper(temperature=0.5    )
         __lowerCamelCase		=  FlaxTemperatureLogitsWarper(temperature=1.3    )

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





         """simple docstring"""



         __lowerCamelCase		=  None
         __lowerCamelCase		=  10
         __lowerCamelCase		=  2

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

         __lowerCamelCase		=  FlaxTopKLogitsWarper(3    )

         __lowerCamelCase		=  top_k_warp(a  ,				a  ,				cur_len=a    )

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

         __lowerCamelCase		=  np.broadcast_to(np.arange(a    )[None, :]  ,				(batch_size, length)    ).copy()
         __lowerCamelCase		=  top_k_warp_safety_check(a  ,				a  ,				cur_len=a    )

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






     def    SCREAMING_SNAKE_CASE__   ( self		:		int    ):





         """simple docstring"""



         __lowerCamelCase		=  None
         __lowerCamelCase		=  10
         __lowerCamelCase		=  2

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

         __lowerCamelCase		=  FlaxTopPLogitsWarper(0.8    )
         __lowerCamelCase		=  np.exp(top_p_warp(a  ,				a  ,				cur_len=a    )    )

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

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

         # make ramp_logits more extreme
         __lowerCamelCase		=  ramp_logits[1] * 1_00.0

         # make sure at least 2 tokens are kept
         __lowerCamelCase		=  FlaxTopPLogitsWarper(0.9  ,				min_tokens_to_keep=2  ,				filter_value=0.0    )
         __lowerCamelCase		=  top_p_warp(a  ,				a  ,				cur_len=a    )

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





         """simple docstring"""



         __lowerCamelCase		=  20
         __lowerCamelCase		=  4
         __lowerCamelCase		=  0

         __lowerCamelCase		=  FlaxMinLengthLogitsProcessor(min_length=10  ,				eos_token_id=a    )

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






     def    SCREAMING_SNAKE_CASE__   ( self		:		Optional[Any]    ):





         """simple docstring"""



         __lowerCamelCase		=  20
         __lowerCamelCase		=  4
         __lowerCamelCase		=  0

         __lowerCamelCase		=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a    )

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






     def    SCREAMING_SNAKE_CASE__   ( self		:		List[str]    ):





         """simple docstring"""



         __lowerCamelCase		=  20
         __lowerCamelCase		=  4
         __lowerCamelCase		=  0
         __lowerCamelCase		=  5

         __lowerCamelCase		=  FlaxForcedEOSTokenLogitsProcessor(max_length=a  ,				eos_token_id=a    )

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






     def    SCREAMING_SNAKE_CASE__   ( self		:		Union[str, Any]    ):





         """simple docstring"""



         __lowerCamelCase		=  4
         __lowerCamelCase		=  10
         __lowerCamelCase		=  15
         __lowerCamelCase		=  2
         __lowerCamelCase		=  1
         __lowerCamelCase		=  15

         # dummy input_ids and scores
         __lowerCamelCase		=  ids_tensor((batch_size, sequence_length)  ,				a    )
         __lowerCamelCase		=  input_ids.copy()

         __lowerCamelCase		=  self._get_uniform_logits(a  ,				a    )
         __lowerCamelCase		=  scores.copy()

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

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

         __lowerCamelCase		=  10

         # no processor list
         __lowerCamelCase		=  temp_dist_warp(a  ,				a  ,				cur_len=a    )
         __lowerCamelCase		=  top_k_warp(a  ,				a  ,				cur_len=a    )
         __lowerCamelCase		=  top_p_warp(a  ,				a  ,				cur_len=a    )
         __lowerCamelCase		=  min_dist_proc(a  ,				a  ,				cur_len=a    )
         __lowerCamelCase		=  bos_dist_proc(a  ,				a  ,				cur_len=a    )
         __lowerCamelCase		=  eos_dist_proc(a  ,				a  ,				cur_len=a    )

         # with processor list
         __lowerCamelCase		=  FlaxLogitsProcessorList(
             [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
         __lowerCamelCase		=  processor(a  ,				a  ,				cur_len=a    )

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

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






     def    SCREAMING_SNAKE_CASE__   ( self		:		str    ):





         """simple docstring"""



         __lowerCamelCase		=  4
         __lowerCamelCase		=  10
         __lowerCamelCase		=  15
         __lowerCamelCase		=  2
         __lowerCamelCase		=  1
         __lowerCamelCase		=  15

         # dummy input_ids and scores
         __lowerCamelCase		=  ids_tensor((batch_size, sequence_length)  ,				a    )
         __lowerCamelCase		=  input_ids.copy()

         __lowerCamelCase		=  self._get_uniform_logits(a  ,				a    )
         __lowerCamelCase		=  scores.copy()

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

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

         __lowerCamelCase		=  10

         # no processor list
         def run_no_processor_list(a		:		Dict  ,				a		:		List[str]  ,				a		:		Dict    ):
             __lowerCamelCase		=  temp_dist_warp(a  ,				a  ,				cur_len=a    )
             __lowerCamelCase		=  top_k_warp(a  ,				a  ,				cur_len=a    )
             __lowerCamelCase		=  top_p_warp(a  ,				a  ,				cur_len=a    )
             __lowerCamelCase		=  min_dist_proc(a  ,				a  ,				cur_len=a    )
             __lowerCamelCase		=  bos_dist_proc(a  ,				a  ,				cur_len=a    )
             __lowerCamelCase		=  eos_dist_proc(a  ,				a  ,				cur_len=a    )
             return scores

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

         __lowerCamelCase		=  jax.jit(a    )
         __lowerCamelCase		=  jax.jit(a    )

         __lowerCamelCase		=  jitted_run_no_processor_list(a  ,				a  ,				a    )
         __lowerCamelCase		=  jitted_run_processor_list(a  ,				a  ,				a    )

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

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








from __future__ import annotations

import collections
import pprint
from pathlib import Path


def __A    (      __lowerCAmelCase	)->  str:





						"""simple docstring"""



						return "".join(sorted(__lowerCAmelCase	)	)


def __A    (      __lowerCAmelCase	)->  list[str]:





						"""simple docstring"""



						return word_by_signature[signature(__lowerCAmelCase	)]


_a      	=  Path(__file__).parent.joinpath('''words.txt''').read_text(encoding='''utf-8''')
_a      	=  sorted({word.strip().lower() for word in data.splitlines()})

_a      	=  collections.defaultdict(list)
for word in word_list:
		word_by_signature[signature(word)].append(word)

if __name__ == "__main__":
		_a      	=  {word: anagram(word) for word in word_list if len(anagram(word)) > 1}

		with open('''anagrams.txt''', '''w''') as file:
				file.write('''all_anagrams = \n ''')
				file.write(pprint.pformat(all_anagrams))
