


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




				def   a__(	self    :				str   ,					lowerCAmelCase    :				int   ,					lowerCAmelCase    :				int   )->  int:

										"""simple docstring"""




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




				def   a__(	self    :				Tuple   )->  List[str]:

										"""simple docstring"""




										UpperCAmelCase         =			None
										UpperCAmelCase         =			20

										UpperCAmelCase         =			self._get_uniform_logits(batch_size=2   ,					length=lowerCAmelCase   )

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

										# compute softmax
										UpperCAmelCase         =			jax.nn.softmax(lowerCAmelCase   ,					axis=-1   )

										UpperCAmelCase         =			FlaxTemperatureLogitsWarper(temperature=0.5   )
										UpperCAmelCase         =			FlaxTemperatureLogitsWarper(temperature=1.3   )

										UpperCAmelCase         =			jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase   ,					scores.copy()   ,					cur_len=lowerCAmelCase   )   ,					axis=-1   )
										UpperCAmelCase         =			jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase   ,					scores.copy()   ,					cur_len=lowerCAmelCase   )   ,					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    :				List[str]   )->  List[str]:

										"""simple docstring"""




										UpperCAmelCase         =			None
										UpperCAmelCase         =			10
										UpperCAmelCase         =			2

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

										UpperCAmelCase         =			FlaxTopKLogitsWarper(3   )

										UpperCAmelCase         =			top_k_warp(lowerCAmelCase   ,					lowerCAmelCase   ,					cur_len=lowerCAmelCase   )

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

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

										# 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    :				Dict   )->  Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase         =			None
										UpperCAmelCase         =			10
										UpperCAmelCase         =			2

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

										UpperCAmelCase         =			FlaxTopPLogitsWarper(0.8   )
										UpperCAmelCase         =			np.exp(top_p_warp(lowerCAmelCase   ,					lowerCAmelCase   ,					cur_len=lowerCAmelCase   )   )

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

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

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

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

										# 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    :				int   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase         =			20
										UpperCAmelCase         =			4
										UpperCAmelCase         =			0

										UpperCAmelCase         =			FlaxMinLengthLogitsProcessor(min_length=10   ,					eos_token_id=lowerCAmelCase   )

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




				def   a__(	self    :				List[Any]   )->  Tuple:

										"""simple docstring"""




										UpperCAmelCase         =			20
										UpperCAmelCase         =			4
										UpperCAmelCase         =			0

										UpperCAmelCase         =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase   )

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




				def   a__(	self    :				Optional[int]   )->  Tuple:

										"""simple docstring"""




										UpperCAmelCase         =			20
										UpperCAmelCase         =			4
										UpperCAmelCase         =			0
										UpperCAmelCase         =			5

										UpperCAmelCase         =			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase   ,					eos_token_id=lowerCAmelCase   )

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




				def   a__(	self    :				List[str]   )->  Tuple:

										"""simple docstring"""




										UpperCAmelCase         =			4
										UpperCAmelCase         =			10
										UpperCAmelCase         =			15
										UpperCAmelCase         =			2
										UpperCAmelCase         =			1
										UpperCAmelCase         =			15

										# dummy input_ids and scores
										UpperCAmelCase         =			ids_tensor((batch_size, sequence_length)   ,					lowerCAmelCase   )
										UpperCAmelCase         =			input_ids.copy()

										UpperCAmelCase         =			self._get_uniform_logits(lowerCAmelCase   ,					lowerCAmelCase   )
										UpperCAmelCase         =			scores.copy()

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

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

										UpperCAmelCase         =			10

										# no processor list
										UpperCAmelCase         =			temp_dist_warp(lowerCAmelCase   ,					lowerCAmelCase   ,					cur_len=lowerCAmelCase   )
										UpperCAmelCase         =			top_k_warp(lowerCAmelCase   ,					lowerCAmelCase   ,					cur_len=lowerCAmelCase   )
										UpperCAmelCase         =			top_p_warp(lowerCAmelCase   ,					lowerCAmelCase   ,					cur_len=lowerCAmelCase   )
										UpperCAmelCase         =			min_dist_proc(lowerCAmelCase   ,					lowerCAmelCase   ,					cur_len=lowerCAmelCase   )
										UpperCAmelCase         =			bos_dist_proc(lowerCAmelCase   ,					lowerCAmelCase   ,					cur_len=lowerCAmelCase   )
										UpperCAmelCase         =			eos_dist_proc(lowerCAmelCase   ,					lowerCAmelCase   ,					cur_len=lowerCAmelCase   )

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

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

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





				def   a__(	self    :				Dict   )->  str:

										"""simple docstring"""




										UpperCAmelCase         =			4
										UpperCAmelCase         =			10
										UpperCAmelCase         =			15
										UpperCAmelCase         =			2
										UpperCAmelCase         =			1
										UpperCAmelCase         =			15

										# dummy input_ids and scores
										UpperCAmelCase         =			ids_tensor((batch_size, sequence_length)   ,					lowerCAmelCase   )
										UpperCAmelCase         =			input_ids.copy()

										UpperCAmelCase         =			self._get_uniform_logits(lowerCAmelCase   ,					lowerCAmelCase   )
										UpperCAmelCase         =			scores.copy()

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

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

										UpperCAmelCase         =			10

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

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

										UpperCAmelCase         =			jax.jit(lowerCAmelCase   )
										UpperCAmelCase         =			jax.jit(lowerCAmelCase   )

										UpperCAmelCase         =			jitted_run_no_processor_list(lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   )
										UpperCAmelCase         =			jitted_run_processor_list(lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   )

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

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









'''simple docstring'''



# coding=utf-8
# Copyright 2020 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# this script dumps information about the environment

import os
import sys

import transformers


_lowercase    :    Union[str, Any]			     =  """3"""

print("""Python version:""", sys.version)
print("""transformers version:""", transformers.__version__)

try:
	import torch

	print("""Torch version:""", torch.__version__)
	print("""Cuda available:""", torch.cuda.is_available())
	print("""Cuda version:""", torch.version.cuda)
	print("""CuDNN version:""", torch.backends.cudnn.version())
	print("""Number of GPUs available:""", torch.cuda.device_count())
	print("""NCCL version:""", torch.cuda.nccl.version())
except ImportError:
	print("""Torch version:""", None)

try:
	import deepspeed

	print("""DeepSpeed version:""", deepspeed.__version__)
except ImportError:
	print("""DeepSpeed version:""", None)

try:
	import tensorflow as tf

	print("""TensorFlow version:""", tf.__version__)
	print("""TF GPUs available:""", bool(tf.config.list_physical_devices("""GPU""")))
	print("""Number of TF GPUs available:""", len(tf.config.list_physical_devices("""GPU""")))
except ImportError:
	print("""TensorFlow version:""", None)





