


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 						a_     (			self     ,   __snake_case     ,   __snake_case       ):
         snake_case											=							jnp.ones((batch_size, length)       ) / length
         return scores
    def 						a_     (			self       ):
         snake_case											=							None
         snake_case											=							2_0

         snake_case											=							self._get_uniform_logits(batch_size=2     ,   length=__snake_case       )

         # 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(__snake_case     ,   axis=-1       )

         snake_case											=							FlaxTemperatureLogitsWarper(temperature=0.5       )
         snake_case											=							FlaxTemperatureLogitsWarper(temperature=1.3       )

         snake_case											=							jax.nn.softmax(temp_dist_warper_sharper(__snake_case     ,   scores.copy()     ,   cur_len=__snake_case       )     ,   axis=-1       )
         snake_case											=							jax.nn.softmax(temp_dist_warper_smoother(__snake_case     ,   scores.copy()     ,   cur_len=__snake_case       )     ,   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       ):
         snake_case											=							None
         snake_case											=							1_0
         snake_case											=							2

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

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

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

         # 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(__snake_case     ,   __snake_case     ,   atol=1E-3       )       )

         # check edge cases with negative and extreme logits
         snake_case											=							np.broadcast_to(np.arange(__snake_case       )[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(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )

         # 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       ):
         snake_case											=							2_0
         snake_case											=							4
         snake_case											=							0

         snake_case											=							FlaxMinLengthLogitsProcessor(min_length=1_0     ,   eos_token_id=__snake_case       )

         # 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(__snake_case     ,   __snake_case       )
         snake_case											=							min_dist_processor(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         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(__snake_case     ,   __snake_case       )
         snake_case											=							1_5
         snake_case											=							min_dist_processor(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         self.assertFalse(jnp.isinf(__snake_case       ).any()       )
    def 						a_     (			self       ):
         snake_case											=							2_0
         snake_case											=							4
         snake_case											=							0

         snake_case											=							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__snake_case       )

         # 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(__snake_case     ,   __snake_case       )
         snake_case											=							logits_processor(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         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(__snake_case     ,   __snake_case       )
         snake_case											=							logits_processor(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         self.assertFalse(jnp.isinf(__snake_case       ).any()       )
    def 						a_     (			self       ):
         snake_case											=							2_0
         snake_case											=							4
         snake_case											=							0
         snake_case											=							5

         snake_case											=							FlaxForcedEOSTokenLogitsProcessor(max_length=__snake_case     ,   eos_token_id=__snake_case       )

         # 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(__snake_case     ,   __snake_case       )
         snake_case											=							logits_processor(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         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(__snake_case     ,   __snake_case       )
         snake_case											=							logits_processor(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         self.assertFalse(jnp.isinf(__snake_case       ).any()       )
    def 						a_     (			self       ):
         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)     ,   __snake_case       )
         snake_case											=							input_ids.copy()

         snake_case											=							self._get_uniform_logits(__snake_case     ,   __snake_case       )
         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=__snake_case       )
         snake_case											=							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__snake_case       )
         snake_case											=							FlaxForcedEOSTokenLogitsProcessor(max_length=__snake_case     ,   eos_token_id=__snake_case       )

         snake_case											=							1_0

         # no processor list
         snake_case											=							temp_dist_warp(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         snake_case											=							top_k_warp(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         snake_case											=							top_p_warp(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         snake_case											=							min_dist_proc(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         snake_case											=							bos_dist_proc(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )
         snake_case											=							eos_dist_proc(__snake_case     ,   __snake_case     ,   cur_len=__snake_case       )

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

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

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




    def 						a_     (			self       ):
         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)     ,   __snake_case       )
         snake_case											=							input_ids.copy()

         snake_case											=							self._get_uniform_logits(__snake_case     ,   __snake_case       )
         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=__snake_case       )
         snake_case											=							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__snake_case       )
         snake_case											=							FlaxForcedEOSTokenLogitsProcessor(max_length=__snake_case     ,   eos_token_id=__snake_case       )

         snake_case											=							1_0

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

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

         snake_case											=							jax.jit(__snake_case       )
         snake_case											=							jax.jit(__snake_case       )

         snake_case											=							jitted_run_no_processor_list(__snake_case     ,   __snake_case     ,   __snake_case       )
         snake_case											=							jitted_run_processor_list(__snake_case     ,   __snake_case     ,   __snake_case       )

         # scores should be equal
         self.assertTrue(jnp.allclose(__snake_case     ,   __snake_case     ,   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)

_SCREAMING_SNAKE_CASE     :					int      =       logging.getLogger()
def      UpperCAmelCase__		():



     """simple docstring"""






     snake_case											=							argparse.ArgumentParser()
     parser.add_argument('''-f'''							)
     snake_case											=							parser.parse_args()
     return args.f

class     A__   (       snake_case__      ):







    """simple docstring"""

    def 						a_     (			self       ):
         snake_case											=							logging.StreamHandler(sys.stdout       )
         logger.addHandler(__snake_case       )
    def 						a_     (			self     ,   __snake_case       ):
         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(__snake_case     ,   '''argv'''     ,   __snake_case       ):
                   snake_case											=							run_glue_deebert.main()
                   for value in result.values():
                        self.assertGreaterEqual(__snake_case     ,   0.666       )




    @slow
    @require_torch_non_multi_gpu
    def 						a_     (			self       ):
         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(__snake_case       )

         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(__snake_case       )

         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(__snake_case       )




