




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




   def        lowerCAmelCase__    ( self    ,							a__    ,							a__							)    ->    Tuple:

          '''simple docstring'''




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




   def        lowerCAmelCase__    ( self							)    ->    int:

          '''simple docstring'''




          snake_case_      =				None
          snake_case_      =				20

          snake_case_      =				self._get_uniform_logits(batch_size=2    ,							length=a__							)

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

          # compute softmax
          snake_case_      =				jax.nn.softmax(a__    ,							axis=-1							)

          snake_case_      =				FlaxTemperatureLogitsWarper(temperature=0.5							)
          snake_case_      =				FlaxTemperatureLogitsWarper(temperature=1.3							)

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

          '''simple docstring'''




          snake_case_      =				None
          snake_case_      =				10
          snake_case_      =				2

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

          snake_case_      =				FlaxTopKLogitsWarper(3							)

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

          snake_case_      =				np.broadcast_to(np.arange(a__							)[None, :]    ,							(batch_size, length)							).copy()
          snake_case_      =				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        lowerCAmelCase__    ( self							)    ->    Optional[Any]:

          '''simple docstring'''




          snake_case_      =				None
          snake_case_      =				10
          snake_case_      =				2

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

          snake_case_      =				FlaxTopPLogitsWarper(0.8							)
          snake_case_      =				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
          snake_case_      =				np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]							)
          self.assertTrue(np.allclose(a__    ,							a__    ,							atol=1e-3							)							)

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

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

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

          '''simple docstring'''




          snake_case_      =				20
          snake_case_      =				4
          snake_case_      =				0

          snake_case_      =				FlaxMinLengthLogitsProcessor(min_length=10    ,							eos_token_id=a__							)

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




   def        lowerCAmelCase__    ( self							)    ->    List[str]:

          '''simple docstring'''




          snake_case_      =				20
          snake_case_      =				4
          snake_case_      =				0

          snake_case_      =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a__							)

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




   def        lowerCAmelCase__    ( self							)    ->    Optional[int]:

          '''simple docstring'''




          snake_case_      =				20
          snake_case_      =				4
          snake_case_      =				0
          snake_case_      =				5

          snake_case_      =				FlaxForcedEOSTokenLogitsProcessor(max_length=a__    ,							eos_token_id=a__							)

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




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

          '''simple docstring'''




          snake_case_      =				4
          snake_case_      =				10
          snake_case_      =				15
          snake_case_      =				2
          snake_case_      =				1
          snake_case_      =				15

          # dummy input_ids and scores
          snake_case_      =				ids_tensor((batch_size, sequence_length)    ,							a__							)
          snake_case_      =				input_ids.copy()

          snake_case_      =				self._get_uniform_logits(a__    ,							a__							)
          snake_case_      =				scores.copy()

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

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

          snake_case_      =				10

          # no processor list
          snake_case_      =				temp_dist_warp(a__    ,							a__    ,							cur_len=a__							)
          snake_case_      =				top_k_warp(a__    ,							a__    ,							cur_len=a__							)
          snake_case_      =				top_p_warp(a__    ,							a__    ,							cur_len=a__							)
          snake_case_      =				min_dist_proc(a__    ,							a__    ,							cur_len=a__							)
          snake_case_      =				bos_dist_proc(a__    ,							a__    ,							cur_len=a__							)
          snake_case_      =				eos_dist_proc(a__    ,							a__    ,							cur_len=a__							)

          # with processor list
          snake_case_      =				FlaxLogitsProcessorList(
              [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
          snake_case_      =				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        lowerCAmelCase__    ( self							)    ->    Tuple:

          '''simple docstring'''




          snake_case_      =				4
          snake_case_      =				10
          snake_case_      =				15
          snake_case_      =				2
          snake_case_      =				1
          snake_case_      =				15

          # dummy input_ids and scores
          snake_case_      =				ids_tensor((batch_size, sequence_length)    ,							a__							)
          snake_case_      =				input_ids.copy()

          snake_case_      =				self._get_uniform_logits(a__    ,							a__							)
          snake_case_      =				scores.copy()

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

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

          snake_case_      =				10

          # no processor list
          def run_no_processor_list(a__    ,							a__    ,							a__							):
                 snake_case_      =				temp_dist_warp(a__    ,							a__    ,							cur_len=a__							)
                 snake_case_      =				top_k_warp(a__    ,							a__    ,							cur_len=a__							)
                 snake_case_      =				top_p_warp(a__    ,							a__    ,							cur_len=a__							)
                 snake_case_      =				min_dist_proc(a__    ,							a__    ,							cur_len=a__							)
                 snake_case_      =				bos_dist_proc(a__    ,							a__    ,							cur_len=a__							)
                 snake_case_      =				eos_dist_proc(a__    ,							a__    ,							cur_len=a__							)
                 return scores

          # with processor list
          def run_processor_list(a__    ,							a__    ,							a__							):
                 snake_case_      =				FlaxLogitsProcessorList(
                     [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
                 snake_case_      =				processor(a__    ,							a__    ,							cur_len=a__							)
                 return scores

          snake_case_      =				jax.jit(a__							)
          snake_case_      =				jax.jit(a__							)

          snake_case_      =				jitted_run_no_processor_list(a__    ,							a__    ,							a__							)
          snake_case_      =				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()							)










'''simple docstring'''





import logging
import os

from .state import PartialState




class 			_snake_case		(     logging.LoggerAdapter							):




			@staticmethod
			def        lowerCAmelCase__    ( a__							)    ->    Optional[Any]:

										'''simple docstring'''




										snake_case_      =				PartialState()
										return not main_process_only or (main_process_only and state.is_main_process)




			def        lowerCAmelCase__    ( self    ,							a__    ,							a__    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										if PartialState._shared_state == {}:
																	raise RuntimeError(
																	    "You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility."							)
										snake_case_      =				kwargs.pop("main_process_only"    ,							a__							)
										snake_case_      =				kwargs.pop("in_order"    ,							a__							)

										if self.isEnabledFor(a__							):
																	if self._should_log(a__							):
																								snake_case_ ,				snake_case_      =				self.process(a__    ,							a__							)
																								self.logger.log(a__    ,							a__    ,							*a__    ,							**a__							)

																	elif in_order:
																								snake_case_      =				PartialState()
																								for i in range(state.num_processes							):
																															if i == state.process_index:
																																						snake_case_ ,				snake_case_      =				self.process(a__    ,							a__							)
																																						self.logger.log(a__    ,							a__    ,							*a__    ,							**a__							)
																															state.wait_for_everyone()


def 					UpperCamelCase_(     snake_case							:	str							,    snake_case							:	str = None					):

							'''simple docstring'''






							if log_level is None:
														snake_case_      =				os.environ.get("ACCELERATE_LOG_LEVEL"							,    snake_case					)
							snake_case_      =				logging.getLogger(snake_case					)
							if log_level is not None:
														logger.setLevel(log_level.upper()					)
														logger.root.setLevel(log_level.upper()					)
							return MultiProcessAdapter(snake_case							,    {}					)




