





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  ):




						"""simple docstring"""







						def 				__snake_case			(		self      ,							a__      ,							a__):



													"""simple docstring"""





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


						def 				__snake_case			(		self):



													"""simple docstring"""





													_lowerCamelCase							:     List[Any]	=		None
													_lowerCamelCase							:     Union[str, Any]	=		20

													_lowerCamelCase							:     List[str]	=		self._get_uniform_logits(batch_size=2      ,							length=_a)

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

													# compute softmax
													_lowerCamelCase							:     str	=		jax.nn.softmax(_a      ,							axis=-1)

													_lowerCamelCase							:     Union[str, Any]	=		FlaxTemperatureLogitsWarper(temperature=0.5)
													_lowerCamelCase							:     Tuple	=		FlaxTemperatureLogitsWarper(temperature=1.3)

													_lowerCamelCase							:     Union[str, Any]	=		jax.nn.softmax(temp_dist_warper_sharper(_a      ,							scores.copy()      ,							cur_len=_a)      ,							axis=-1)
													_lowerCamelCase							:     Optional[int]	=		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 				__snake_case			(		self):



													"""simple docstring"""





													_lowerCamelCase							:     int	=		None
													_lowerCamelCase							:     List[Any]	=		10
													_lowerCamelCase							:     int	=		2

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

													_lowerCamelCase							:     int	=		FlaxTopKLogitsWarper(3)

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

													_lowerCamelCase							:     Any	=		np.broadcast_to(np.arange(_a)[None, :]      ,							(batch_size, length)).copy()
													_lowerCamelCase							:     Dict	=		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 				__snake_case			(		self):



													"""simple docstring"""





													_lowerCamelCase							:     Union[str, Any]	=		None
													_lowerCamelCase							:     Optional[int]	=		10
													_lowerCamelCase							:     Optional[Any]	=		2

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

													_lowerCamelCase							:     Any	=		FlaxTopPLogitsWarper(0.8)
													_lowerCamelCase							:     Union[str, Any]	=		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							:     List[Any]	=		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							:     List[str]	=		np.broadcast_to(np.arange(_a)[None, :]      ,							(batch_size, vocab_size)).copy() - (
													    vocab_size // 2
													)

													# make ramp_logits more extreme
													_lowerCamelCase							:     List[Any]	=		ramp_logits[1] * 100.0

													# make sure at least 2 tokens are kept
													_lowerCamelCase							:     Tuple	=		FlaxTopPLogitsWarper(0.9      ,							min_tokens_to_keep=2      ,							filter_value=0.0)
													_lowerCamelCase							:     int	=		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 				__snake_case			(		self):



													"""simple docstring"""





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

													_lowerCamelCase							:     str	=		FlaxMinLengthLogitsProcessor(min_length=10      ,							eos_token_id=_a)

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


						def 				__snake_case			(		self):



													"""simple docstring"""





													_lowerCamelCase							:     Optional[Any]	=		20
													_lowerCamelCase							:     Any	=		4
													_lowerCamelCase							:     str	=		0

													_lowerCamelCase							:     Any	=		FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_a)

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


						def 				__snake_case			(		self):



													"""simple docstring"""





													_lowerCamelCase							:     Optional[int]	=		20
													_lowerCamelCase							:     Optional[int]	=		4
													_lowerCamelCase							:     Any	=		0
													_lowerCamelCase							:     Dict	=		5

													_lowerCamelCase							:     List[Any]	=		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							:     List[Any]	=		ids_tensor((batch_size, 4)      ,							vocab_size=20)
													_lowerCamelCase							:     Optional[Any]	=		4
													_lowerCamelCase							:     str	=		self._get_uniform_logits(_a      ,							_a)
													_lowerCamelCase							:     Any	=		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							:     Any	=		3
													_lowerCamelCase							:     int	=		self._get_uniform_logits(_a      ,							_a)
													_lowerCamelCase							:     List[str]	=		logits_processor(_a      ,							_a      ,							cur_len=_a)
													self.assertFalse(jnp.isinf(_a).any())


						def 				__snake_case			(		self):



													"""simple docstring"""





													_lowerCamelCase							:     Dict	=		4
													_lowerCamelCase							:     Tuple	=		10
													_lowerCamelCase							:     int	=		15
													_lowerCamelCase							:     int	=		2
													_lowerCamelCase							:     Any	=		1
													_lowerCamelCase							:     Tuple	=		15

													# dummy input_ids and scores
													_lowerCamelCase							:     List[Any]	=		ids_tensor((batch_size, sequence_length)      ,							_a)
													_lowerCamelCase							:     Optional[int]	=		input_ids.copy()

													_lowerCamelCase							:     Union[str, Any]	=		self._get_uniform_logits(_a      ,							_a)
													_lowerCamelCase							:     Optional[int]	=		scores.copy()

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

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

													_lowerCamelCase							:     Union[str, Any]	=		10

													# no processor list
													_lowerCamelCase							:     Optional[int]	=		temp_dist_warp(_a      ,							_a      ,							cur_len=_a)
													_lowerCamelCase							:     Tuple	=		top_k_warp(_a      ,							_a      ,							cur_len=_a)
													_lowerCamelCase							:     List[str]	=		top_p_warp(_a      ,							_a      ,							cur_len=_a)
													_lowerCamelCase							:     str	=		min_dist_proc(_a      ,							_a      ,							cur_len=_a)
													_lowerCamelCase							:     Union[str, Any]	=		bos_dist_proc(_a      ,							_a      ,							cur_len=_a)
													_lowerCamelCase							:     Optional[int]	=		eos_dist_proc(_a      ,							_a      ,							cur_len=_a)

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



													"""simple docstring"""





													_lowerCamelCase							:     int	=		4
													_lowerCamelCase							:     Tuple	=		10
													_lowerCamelCase							:     Optional[Any]	=		15
													_lowerCamelCase							:     str	=		2
													_lowerCamelCase							:     str	=		1
													_lowerCamelCase							:     Dict	=		15

													# dummy input_ids and scores
													_lowerCamelCase							:     int	=		ids_tensor((batch_size, sequence_length)      ,							_a)
													_lowerCamelCase							:     List[Any]	=		input_ids.copy()

													_lowerCamelCase							:     int	=		self._get_uniform_logits(_a      ,							_a)
													_lowerCamelCase							:     Optional[int]	=		scores.copy()

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

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

													_lowerCamelCase							:     List[str]	=		10

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

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

													_lowerCamelCase							:     int	=		jax.jit(_a)
													_lowerCamelCase							:     str	=		jax.jit(_a)

													_lowerCamelCase							:     Tuple	=		jitted_run_no_processor_list(_a      ,							_a      ,							_a)
													_lowerCamelCase							:     str	=		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())








import importlib.util
import os
import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import (
    is_accelerate_available,
    is_flax_available,
    is_safetensors_available,
    is_tf_available,
    is_torch_available,
)
from . import BaseTransformersCLICommand




def    __UpperCAmelCase(      lowercase_      ):
       return EnvironmentCommand()




def    __UpperCAmelCase(      lowercase_      ):
       return EnvironmentCommand(args.accelerate_config_file      )







class 			__A       (      lowerCamelCase__  ):




      """simple docstring"""







      @staticmethod
      def 				__snake_case			(		a__):



             """simple docstring"""





             _lowerCamelCase							:     List[Any]	=		parser.add_parser('''env''')
             download_parser.set_defaults(func=a__)
             download_parser.add_argument(
                 '''--accelerate-config_file'''      ,							default=a__      ,							help='''The accelerate config file to use for the default values in the launching script.'''      ,							)
             download_parser.set_defaults(func=a__)


      def __init__(		self      ,							a__      ,							*a__):



             """simple docstring"""





             _lowerCamelCase							:     str	=		accelerate_config_file


      def 				__snake_case			(		self):



             """simple docstring"""





             _lowerCamelCase							:     Union[str, Any]	=		'''not installed'''
             if is_safetensors_available():
                    import safetensors

                    _lowerCamelCase							:     Optional[Any]	=		safetensors.__version__
             elif importlib.util.find_spec('''safetensors''') is not None:
                    import safetensors

                    _lowerCamelCase							:     Optional[int]	=		F"""{safetensors.__version__} but is ignored because of PyTorch version too old."""

             _lowerCamelCase							:     Union[str, Any]	=		'''not installed'''
             _lowerCamelCase							:     Any	=		'''not found'''
             if is_accelerate_available():
                    import accelerate
                    from accelerate.commands.config import default_config_file, load_config_from_file

                    _lowerCamelCase							:     Optional[int]	=		accelerate.__version__
                    # Get the default from the config file.
                    if self._accelerate_config_file is not None or os.path.isfile(a__):
                           _lowerCamelCase							:     Optional[int]	=		load_config_from_file(self._accelerate_config_file).to_dict()

                    _lowerCamelCase							:     str	=		(
                        '''\n'''.join([F"""\t- {prop}: {val}""" for prop, val in accelerate_config.items()])
                        if isinstance(a__      ,							a__)
                        else F"""\t{accelerate_config}"""
                    )

             _lowerCamelCase							:     List[Any]	=		'''not installed'''
             _lowerCamelCase							:     Tuple	=		'''NA'''
             if is_torch_available():
                    import torch

                    _lowerCamelCase							:     int	=		torch.__version__
                    _lowerCamelCase							:     List[str]	=		torch.cuda.is_available()

             _lowerCamelCase							:     str	=		'''not installed'''
             _lowerCamelCase							:     Union[str, Any]	=		'''NA'''
             if is_tf_available():
                    import tensorflow as tf

                    _lowerCamelCase							:     List[str]	=		tf.__version__
                    try:
                           # deprecated in v2.1
                           _lowerCamelCase							:     Optional[int]	=		tf.test.is_gpu_available()
                    except AttributeError:
                           # returns list of devices, convert to bool
                           _lowerCamelCase							:     Optional[int]	=		bool(tf.config.list_physical_devices('''GPU'''))

             _lowerCamelCase							:     str	=		'''not installed'''
             _lowerCamelCase							:     List[Any]	=		'''not installed'''
             _lowerCamelCase							:     List[Any]	=		'''not installed'''
             _lowerCamelCase							:     Optional[int]	=		'''NA'''
             if is_flax_available():
                    import flax
                    import jax
                    import jaxlib

                    _lowerCamelCase							:     Any	=		flax.__version__
                    _lowerCamelCase							:     str	=		jax.__version__
                    _lowerCamelCase							:     Any	=		jaxlib.__version__
                    _lowerCamelCase							:     int	=		jax.lib.xla_bridge.get_backend().platform

             _lowerCamelCase							:     int	=		{
                 '''`transformers` version''': version,
                 '''Platform''': platform.platform(),
                 '''Python version''': platform.python_version(),
                 '''Huggingface_hub version''': huggingface_hub.__version__,
                 '''Safetensors version''': F"""{safetensors_version}""",
                 '''Accelerate version''': F"""{accelerate_version}""",
                 '''Accelerate config''': F"""{accelerate_config_str}""",
                 '''PyTorch version (GPU?)''': F"""{pt_version} ({pt_cuda_available})""",
                 '''Tensorflow version (GPU?)''': F"""{tf_version} ({tf_cuda_available})""",
                 '''Flax version (CPU?/GPU?/TPU?)''': F"""{flax_version} ({jax_backend})""",
                 '''Jax version''': F"""{jax_version}""",
                 '''JaxLib version''': F"""{jaxlib_version}""",
                 '''Using GPU in script?''': '''<fill in>''',
                 '''Using distributed or parallel set-up in script?''': '''<fill in>''',
             }

             print('''\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n''')
             print(self.format_dict(a__))

             return info



      @staticmethod
      def 				__snake_case			(		a__):



             """simple docstring"""





             return "\n".join([F"""- {prop}: {val}""" for prop, val in d.items()]) + "\n"

