






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






    def    UpperCamelCase       (       self     :							List[Any]       ,						__SCREAMING_SNAKE_CASE     :							int       ,						__SCREAMING_SNAKE_CASE     :							int					)     ->		Tuple:
      lowerCamelCase_         =						jnp.ones((batch_size, length)					) / length
      return scores






    def    UpperCamelCase       (       self     :							Optional[Any]					)     ->		int:
      lowerCamelCase_         =						None
      lowerCamelCase_         =						20

      lowerCamelCase_         =						self._get_uniform_logits(batch_size=2       ,						length=__SCREAMING_SNAKE_CASE					)

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

      # compute softmax
      lowerCamelCase_         =						jax.nn.softmax(__SCREAMING_SNAKE_CASE       ,						axis=-1					)

      lowerCamelCase_         =						FlaxTemperatureLogitsWarper(temperature=0.5					)
      lowerCamelCase_         =						FlaxTemperatureLogitsWarper(temperature=1.3					)

      lowerCamelCase_         =						jax.nn.softmax(temp_dist_warper_sharper(__SCREAMING_SNAKE_CASE       ,						scores.copy()       ,						cur_len=__SCREAMING_SNAKE_CASE					)       ,						axis=-1					)
      lowerCamelCase_         =						jax.nn.softmax(temp_dist_warper_smoother(__SCREAMING_SNAKE_CASE       ,						scores.copy()       ,						cur_len=__SCREAMING_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    UpperCamelCase       (       self     :							Tuple					)     ->		Any:
      lowerCamelCase_         =						None
      lowerCamelCase_         =						10
      lowerCamelCase_         =						2

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

      lowerCamelCase_         =						FlaxTopKLogitsWarper(3					)

      lowerCamelCase_         =						top_k_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_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
      lowerCamelCase_         =						5
      lowerCamelCase_         =						FlaxTopKLogitsWarper(top_k=1       ,						filter_value=0.0       ,						min_tokens_to_keep=3					)

      lowerCamelCase_         =						np.broadcast_to(np.arange(__SCREAMING_SNAKE_CASE					)[None, :]       ,						(batch_size, length)					).copy()
      lowerCamelCase_         =						top_k_warp_safety_check(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_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    UpperCamelCase       (       self     :							Optional[Any]					)     ->		Optional[int]:
      lowerCamelCase_         =						None
      lowerCamelCase_         =						10
      lowerCamelCase_         =						2

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

      lowerCamelCase_         =						FlaxTopPLogitsWarper(0.8					)
      lowerCamelCase_         =						np.exp(top_p_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)					)

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

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

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

      # make sure at least 2 tokens are kept
      lowerCamelCase_         =						FlaxTopPLogitsWarper(0.9       ,						min_tokens_to_keep=2       ,						filter_value=0.0					)
      lowerCamelCase_         =						top_p_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_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    UpperCamelCase       (       self     :							List[Any]					)     ->		Tuple:
      lowerCamelCase_         =						20
      lowerCamelCase_         =						4
      lowerCamelCase_         =						0

      lowerCamelCase_         =						FlaxMinLengthLogitsProcessor(min_length=10       ,						eos_token_id=__SCREAMING_SNAKE_CASE					)

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






    def    UpperCamelCase       (       self     :							str					)     ->		List[str]:
      lowerCamelCase_         =						20
      lowerCamelCase_         =						4
      lowerCamelCase_         =						0

      lowerCamelCase_         =						FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE					)

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






    def    UpperCamelCase       (       self     :							Tuple					)     ->		Optional[int]:
      lowerCamelCase_         =						20
      lowerCamelCase_         =						4
      lowerCamelCase_         =						0
      lowerCamelCase_         =						5

      lowerCamelCase_         =						FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE       ,						eos_token_id=__SCREAMING_SNAKE_CASE					)

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






    def    UpperCamelCase       (       self     :							Dict					)     ->		Dict:
      lowerCamelCase_         =						4
      lowerCamelCase_         =						10
      lowerCamelCase_         =						15
      lowerCamelCase_         =						2
      lowerCamelCase_         =						1
      lowerCamelCase_         =						15

      # dummy input_ids and scores
      lowerCamelCase_         =						ids_tensor((batch_size, sequence_length)       ,						__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						input_ids.copy()

      lowerCamelCase_         =						self._get_uniform_logits(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						scores.copy()

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

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

      lowerCamelCase_         =						10

      # no processor list
      lowerCamelCase_         =						temp_dist_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						top_k_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						top_p_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						min_dist_proc(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						bos_dist_proc(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						eos_dist_proc(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)

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

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

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






    def    UpperCamelCase       (       self     :							List[Any]					)     ->		str:
      lowerCamelCase_         =						4
      lowerCamelCase_         =						10
      lowerCamelCase_         =						15
      lowerCamelCase_         =						2
      lowerCamelCase_         =						1
      lowerCamelCase_         =						15

      # dummy input_ids and scores
      lowerCamelCase_         =						ids_tensor((batch_size, sequence_length)       ,						__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						input_ids.copy()

      lowerCamelCase_         =						self._get_uniform_logits(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						scores.copy()

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

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

      lowerCamelCase_         =						10

      # no processor list
      def run_no_processor_list(__SCREAMING_SNAKE_CASE     :							Tuple       ,						__SCREAMING_SNAKE_CASE     :							str       ,						__SCREAMING_SNAKE_CASE     :							List[str]					):
        lowerCamelCase_         =						temp_dist_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
        lowerCamelCase_         =						top_k_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
        lowerCamelCase_         =						top_p_warp(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
        lowerCamelCase_         =						min_dist_proc(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
        lowerCamelCase_         =						bos_dist_proc(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
        lowerCamelCase_         =						eos_dist_proc(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						cur_len=__SCREAMING_SNAKE_CASE					)
        return scores

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

      lowerCamelCase_         =						jax.jit(__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						jax.jit(__SCREAMING_SNAKE_CASE					)

      lowerCamelCase_         =						jitted_run_no_processor_list(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						jitted_run_processor_list(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)

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

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












"""simple docstring"""





import argparse
import json
import subprocess







def      lowerCamelCase__       (							_lowerCamelCase       :		Tuple		,				_lowerCamelCase       :		str )      ->				List[Any]:
  lowerCamelCase_         =						[]

  lowerCamelCase_         =						(
      F'''curl -H "Accept: application/vnd.github+json" -H "Authorization: Bearer {token}"'''
      ' https://api.github.com/repos/huggingface/transformers/actions/runners'
  )
  lowerCamelCase_         =						subprocess.run(_lowerCamelCase		,				shell=_lowerCamelCase		,				stdout=subprocess.PIPE )
  lowerCamelCase_         =						output.stdout.decode('utf-8' )
  lowerCamelCase_         =						json.loads(_lowerCamelCase )

  lowerCamelCase_         =						status['runners']
  for runner in runners:
    if runner["name"] in target_runners:
      if runner["status"] == "offline":
        offline_runners.append(_lowerCamelCase )

    # save the result so we can report them on Slack
  with open('offline_runners.txt'		,				'w' ) as fp:
    fp.write(json.dumps(_lowerCamelCase ) )

  if len(_lowerCamelCase ) > 0:
    lowerCamelCase_         =						'\n'.join([x['name'] for x in offline_runners] )
    raise ValueError(F'''The following runners are offline:\n{failed}''' )


if __name__ == "__main__":







  def      lowerCamelCase__       (							_lowerCamelCase       :		Dict )      ->				Tuple:
    return values.split(',' )

  _SCREAMING_SNAKE_CASE					:				Union[str, Any]					    =						argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '''--target_runners''',
      default=None,
      type=list_str,
      required=True,
      help='''Comma-separated list of runners to check status.''',
  )

  parser.add_argument(
      '''--token''', default=None, type=str, required=True, help='''A token that has actions:read permission.'''
  )
  _SCREAMING_SNAKE_CASE					:				Any					    =						parser.parse_args()






  get_runner_status(args.target_runners, args.token)




