






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



				"""simple docstring"""



				def 					__magic_name__							(      self    :    List[Any]			,		UpperCamelCase__    :    int			,		UpperCamelCase__    :    int):






											'''simple docstring'''







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



				def 					__magic_name__							(      self    :    Any):






											'''simple docstring'''







											snake_case__	  =       None
											snake_case__	  =       2_0

											snake_case__	  =       self._get_uniform_logits(batch_size=2			,		length=UpperCamelCase__)

											# 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, 1_0].set((1 / length) - 0.4)  # valley, 1st batch

											# compute softmax
											snake_case__	  =       jax.nn.softmax(UpperCamelCase__			,		axis=-1)

											snake_case__	  =       FlaxTemperatureLogitsWarper(temperature=0.5)
											snake_case__	  =       FlaxTemperatureLogitsWarper(temperature=1.3)

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






											'''simple docstring'''







											snake_case__	  =       None
											snake_case__	  =       1_0
											snake_case__	  =       2

											# create ramp distribution
											snake_case__	  =       np.broadcast_to(np.arange(UpperCamelCase__)[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(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)

											# 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(UpperCamelCase__)[None, :]			,		(batch_size, length)).copy()
											snake_case__	  =       top_k_warp_safety_check(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)

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



				def 					__magic_name__							(      self    :    Any):






											'''simple docstring'''







											snake_case__	  =       None
											snake_case__	  =       1_0
											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.15, 0.3, 0.3, 0.25]]))

											snake_case__	  =       FlaxTopPLogitsWarper(0.8)
											snake_case__	  =       np.exp(top_p_warp(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__))

											# 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.25]])
											self.assertTrue(np.allclose(UpperCamelCase__			,		UpperCamelCase__			,		atol=1E-3))

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

											# make ramp_logits more extreme
											snake_case__	  =       ramp_logits[1] * 1_00.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(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)

											# 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 					__magic_name__							(      self    :    Optional[Any]):






											'''simple docstring'''







											snake_case__	  =       2_0
											snake_case__	  =       4
											snake_case__	  =       0

											snake_case__	  =       FlaxMinLengthLogitsProcessor(min_length=1_0			,		eos_token_id=UpperCamelCase__)

											# check that min length is applied at length 5
											snake_case__	  =       ids_tensor((batch_size, 2_0)			,		vocab_size=2_0)
											snake_case__	  =       5
											snake_case__	  =       self._get_uniform_logits(UpperCamelCase__			,		UpperCamelCase__)
											snake_case__	  =       min_dist_processor(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											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(UpperCamelCase__			,		UpperCamelCase__)
											snake_case__	  =       1_5
											snake_case__	  =       min_dist_processor(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											self.assertFalse(jnp.isinf(UpperCamelCase__).any())



				def 					__magic_name__							(      self    :    List[str]):






											'''simple docstring'''







											snake_case__	  =       2_0
											snake_case__	  =       4
											snake_case__	  =       0

											snake_case__	  =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__)

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



				def 					__magic_name__							(      self    :    Optional[Any]):






											'''simple docstring'''







											snake_case__	  =       2_0
											snake_case__	  =       4
											snake_case__	  =       0
											snake_case__	  =       5

											snake_case__	  =       FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__			,		eos_token_id=UpperCamelCase__)

											# 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=2_0)
											snake_case__	  =       4
											snake_case__	  =       self._get_uniform_logits(UpperCamelCase__			,		UpperCamelCase__)
											snake_case__	  =       logits_processor(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											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(UpperCamelCase__			,		UpperCamelCase__)
											snake_case__	  =       logits_processor(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											self.assertFalse(jnp.isinf(UpperCamelCase__).any())



				def 					__magic_name__							(      self    :    Tuple):






											'''simple docstring'''







											snake_case__	  =       4
											snake_case__	  =       1_0
											snake_case__	  =       1_5
											snake_case__	  =       2
											snake_case__	  =       1
											snake_case__	  =       1_5

											# dummy input_ids and scores
											snake_case__	  =       ids_tensor((batch_size, sequence_length)			,		UpperCamelCase__)
											snake_case__	  =       input_ids.copy()

											snake_case__	  =       self._get_uniform_logits(UpperCamelCase__			,		UpperCamelCase__)
											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=1_0			,		eos_token_id=UpperCamelCase__)
											snake_case__	  =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__)
											snake_case__	  =       FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__			,		eos_token_id=UpperCamelCase__)

											snake_case__	  =       1_0

											# no processor list
											snake_case__	  =       temp_dist_warp(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											snake_case__	  =       top_k_warp(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											snake_case__	  =       top_p_warp(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											snake_case__	  =       min_dist_proc(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											snake_case__	  =       bos_dist_proc(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
											snake_case__	  =       eos_dist_proc(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)

											# 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(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)

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

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



				def 					__magic_name__							(      self    :    Optional[int]):






											'''simple docstring'''







											snake_case__	  =       4
											snake_case__	  =       1_0
											snake_case__	  =       1_5
											snake_case__	  =       2
											snake_case__	  =       1
											snake_case__	  =       1_5

											# dummy input_ids and scores
											snake_case__	  =       ids_tensor((batch_size, sequence_length)			,		UpperCamelCase__)
											snake_case__	  =       input_ids.copy()

											snake_case__	  =       self._get_uniform_logits(UpperCamelCase__			,		UpperCamelCase__)
											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=1_0			,		eos_token_id=UpperCamelCase__)
											snake_case__	  =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__)
											snake_case__	  =       FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__			,		eos_token_id=UpperCamelCase__)

											snake_case__	  =       1_0

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

											# with processor list
											def run_processor_list(UpperCamelCase__    :    Dict			,		UpperCamelCase__    :    Optional[Any]			,		UpperCamelCase__    :    str):
																		snake_case__	  =       FlaxLogitsProcessorList(
																		    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
																		snake_case__	  =       processor(UpperCamelCase__			,		UpperCamelCase__			,		cur_len=UpperCamelCase__)
																		return scores

											snake_case__	  =       jax.jit(UpperCamelCase__)
											snake_case__	  =       jax.jit(UpperCamelCase__)

											snake_case__	  =       jitted_run_no_processor_list(UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__)
											snake_case__	  =       jitted_run_processor_list(UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__)

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

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








import argparse
import logging
import sys
from unittest.mock import patch

import run_glue_deebert

from transformers.testing_utils import TestCasePlus, get_gpu_count, require_torch_non_multi_gpu, slow


logging.basicConfig(level=logging.DEBUG)

a__       =     logging.getLogger()

def _UpperCAmelCase   (	):
							snake_case__	  =       argparse.ArgumentParser()
							parser.add_argument("""-f""" )
							snake_case__	  =       parser.parse_args()
							return args.f



class    _lowerCAmelCase			(		lowercase_      ):



				"""simple docstring"""



				def 					__magic_name__							(      self    :    int):






											'''simple docstring'''







											snake_case__	  =       logging.StreamHandler(sys.stdout)
											logger.addHandler(UpperCamelCase__)



				def 					__magic_name__							(      self    :    Union[str, Any]			,		UpperCamelCase__    :    List[str]):






											'''simple docstring'''







											snake_case__	  =       get_gpu_count()

											if n_gpu > 1:
																		pass
																		# XXX: doesn't quite work with n_gpu > 1 https://github.com/huggingface/transformers/issues/10560
																		# script = f"{self.examples_dir_str}/research_projects/deebert/run_glue_deebert.py"
																		# distributed_args = f"-m torch.distributed.launch --nproc_per_node={n_gpu} {script}".split()
																		# cmd = [sys.executable] + distributed_args + args
																		# execute_subprocess_async(cmd, env=self.get_env())
																		# XXX: test the results - need to save them first into .json file
											else:
																		args.insert(0			,		"""run_glue_deebert.py""")
																		with patch.object(UpperCamelCase__			,		"""argv"""			,		UpperCamelCase__):
																									snake_case__	  =       run_glue_deebert.main()
																									for value in result.values():
																																self.assertGreaterEqual(UpperCamelCase__			,		0.6_66)



				@slow
				@require_torch_non_multi_gpu
				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											snake_case__	  =       """
            --model_type roberta
            --model_name_or_path roberta-base
            --task_name MRPC
            --do_train
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --max_seq_length 128
            --per_gpu_eval_batch_size=1
            --per_gpu_train_batch_size=8
            --learning_rate 2e-4
            --num_train_epochs 3
            --overwrite_output_dir
            --seed 42
            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --save_steps 0
            --overwrite_cache
            --eval_after_first_stage
            """.split()
											self.run_and_check(UpperCamelCase__)

											snake_case__	  =       """
            --model_type roberta
            --model_name_or_path ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --task_name MRPC
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --max_seq_length 128
            --eval_each_highway
            --eval_highway
            --overwrite_cache
            --per_gpu_eval_batch_size=1
            """.split()
											self.run_and_check(UpperCamelCase__)

											snake_case__	  =       """
            --model_type roberta
            --model_name_or_path ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --task_name MRPC
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --max_seq_length 128
            --early_exit_entropy 0.1
            --eval_highway
            --overwrite_cache
            --per_gpu_eval_batch_size=1
            """.split()
											self.run_and_check(UpperCamelCase__)
