





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  :       List[str]				,      __lowerCAmelCase  :       int				,      __lowerCAmelCase  :       int  )      ->      List[Any]:







					"""simple docstring"""






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






			def   a_			(     self  :       Optional[Any]  )      ->      str:







					"""simple docstring"""






					A__     =  None
					A__     =  20

					A__     =  self._get_uniform_logits(batch_size=2				,      length=__lowerCamelCase  )

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

					# compute softmax
					A__     =  jax.nn.softmax(__lowerCamelCase				,      axis=-1  )

					A__     =  FlaxTemperatureLogitsWarper(temperature=0.5  )
					A__     =  FlaxTemperatureLogitsWarper(temperature=1.3  )

					A__     =  jax.nn.softmax(temp_dist_warper_sharper(__lowerCamelCase				,      scores.copy()				,      cur_len=__lowerCamelCase  )				,      axis=-1  )
					A__     =  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  :       Optional[int]  )      ->      Optional[int]:







					"""simple docstring"""






					A__     =  None
					A__     =  10
					A__     =  2

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

					A__     =  FlaxTopKLogitsWarper(3  )

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

					A__     =  np.broadcast_to(np.arange(__lowerCamelCase  )[None, :]				,      (batch_size, length)  ).copy()
					A__     =  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  :       int  )      ->      Tuple:







					"""simple docstring"""






					A__     =  None
					A__     =  10
					A__     =  2

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

					A__     =  FlaxTopPLogitsWarper(0.8  )
					A__     =  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
					A__     =  np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]  )
					self.assertTrue(np.allclose(__lowerCamelCase				,      __lowerCamelCase				,      atol=1e-3  )  )

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

					# make ramp_logits more extreme
					A__     =  ramp_logits[1] * 1_0_0.0

					# make sure at least 2 tokens are kept
					A__     =  FlaxTopPLogitsWarper(0.9				,      min_tokens_to_keep=2				,      filter_value=0.0  )
					A__     =  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  :       Any  )      ->      Optional[int]:







					"""simple docstring"""






					A__     =  20
					A__     =  4
					A__     =  0

					A__     =  FlaxMinLengthLogitsProcessor(min_length=10				,      eos_token_id=__lowerCamelCase  )

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






			def   a_			(     self  :       int  )      ->      str:







					"""simple docstring"""






					A__     =  20
					A__     =  4
					A__     =  0

					A__     =  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__lowerCamelCase  )

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






			def   a_			(     self  :       Tuple  )      ->      Dict:







					"""simple docstring"""






					A__     =  20
					A__     =  4
					A__     =  0
					A__     =  5

					A__     =  FlaxForcedEOSTokenLogitsProcessor(max_length=__lowerCamelCase				,      eos_token_id=__lowerCamelCase  )

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






			def   a_			(     self  :       Optional[Any]  )      ->      Union[str, Any]:







					"""simple docstring"""






					A__     =  4
					A__     =  10
					A__     =  15
					A__     =  2
					A__     =  1
					A__     =  15

					# dummy input_ids and scores
					A__     =  ids_tensor((batch_size, sequence_length)				,      __lowerCamelCase  )
					A__     =  input_ids.copy()

					A__     =  self._get_uniform_logits(__lowerCamelCase				,      __lowerCamelCase  )
					A__     =  scores.copy()

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

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

					A__     =  10

					# no processor list
					A__     =  temp_dist_warp(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
					A__     =  top_k_warp(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
					A__     =  top_p_warp(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
					A__     =  min_dist_proc(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
					A__     =  bos_dist_proc(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
					A__     =  eos_dist_proc(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )

					# with processor list
					A__     =  FlaxLogitsProcessorList(
					    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
					A__     =  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  :       List[Any]  )      ->      Optional[Any]:







					"""simple docstring"""






					A__     =  4
					A__     =  10
					A__     =  15
					A__     =  2
					A__     =  1
					A__     =  15

					# dummy input_ids and scores
					A__     =  ids_tensor((batch_size, sequence_length)				,      __lowerCamelCase  )
					A__     =  input_ids.copy()

					A__     =  self._get_uniform_logits(__lowerCamelCase				,      __lowerCamelCase  )
					A__     =  scores.copy()

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

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

					A__     =  10

					# no processor list
					def run_no_processor_list(__lowerCAmelCase  :       List[str]				,      __lowerCAmelCase  :       str				,      __lowerCAmelCase  :       List[Any]  ):
							A__     =  temp_dist_warp(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
							A__     =  top_k_warp(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
							A__     =  top_p_warp(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
							A__     =  min_dist_proc(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
							A__     =  bos_dist_proc(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
							A__     =  eos_dist_proc(__lowerCamelCase				,      __lowerCamelCase				,      cur_len=__lowerCamelCase  )
							return scores

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

					A__     =  jax.jit(__lowerCamelCase  )
					A__     =  jax.jit(__lowerCamelCase  )

					A__     =  jitted_run_no_processor_list(__lowerCamelCase				,      __lowerCamelCase				,      __lowerCamelCase  )
					A__     =  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()  )









import argparse
import json
import os
import time
import zipfile

from get_ci_error_statistics import download_artifact, get_artifacts_links

from transformers import logging


__magic_name__  :  List[Any]					=       logging.get_logger(__name__)
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase				):
  lowerCAmelCase__										=      set()
  lowerCAmelCase__										=      []

  def parse_line(__lowerCAmelCase				):
    for line in fp:
      if isinstance(__lowerCAmelCase   , __lowerCAmelCase				):
        lowerCAmelCase__										=      line.decode('''UTF-8'''				)
      if "warnings summary (final)" in line:
        continue
      # This means we are outside the body of a warning
      elif not line.startswith(''' '''				):
        # process a single warning and move it to `selected_warnings`.
        if len(__lowerCAmelCase				) > 0:
          lowerCAmelCase__										=      '''\n'''.join(__lowerCAmelCase				)
          # Only keep the warnings specified in `targets`
          if any(F""": {x}: """ in warning for x in targets				):
            selected_warnings.add(__lowerCAmelCase				)
          buffer.clear()
        continue
      else:
        lowerCAmelCase__										=      line.strip()
        buffer.append(__lowerCAmelCase				)

  if from_gh:
    for filename in os.listdir(__lowerCAmelCase				):
      lowerCAmelCase__										=      os.path.join(__lowerCAmelCase   , __lowerCAmelCase				)
      if not os.path.isdir(__lowerCAmelCase				):
        # read the file
        if filename != "warnings.txt":
          continue
        with open(__lowerCAmelCase				) as fp:
          parse_line(__lowerCAmelCase				)
  else:
    try:
      with zipfile.ZipFile(__lowerCAmelCase				) as z:
        for filename in z.namelist():
          if not os.path.isdir(__lowerCAmelCase				):
            # read the file
            if filename != "warnings.txt":
              continue
            with z.open(__lowerCAmelCase				) as fp:
              parse_line(__lowerCAmelCase				)
    except Exception:
      logger.warning(
          F"""{artifact_path} is either an invalid zip file or something else wrong. This file is skipped."""				)

  return selected_warnings
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase				):
  lowerCAmelCase__										=      set()

  lowerCAmelCase__										=      [os.path.join(__lowerCAmelCase   , __lowerCAmelCase				) for p in os.listdir(__lowerCAmelCase				) if (p.endswith('''.zip'''				) or from_gh)]
  for p in paths:
    selected_warnings.update(extract_warnings_from_single_artifact(__lowerCAmelCase   , __lowerCAmelCase				)				)

  return selected_warnings


if __name__ == "__main__":
      def 		a_							( __lowerCAmelCase				):
        return values.split(''','''				)

      __magic_name__  :  Optional[int]					=       argparse.ArgumentParser()
      # Required parameters
      parser.add_argument("""--workflow_run_id""", type=str, required=True, help="""A GitHub Actions workflow run id.""")
      parser.add_argument(
          """--output_dir""",
          type=str,
          required=True,
          help="""Where to store the downloaded artifacts and other result files.""",
      )
      parser.add_argument("""--token""", default=None, type=str, help="""A token that has actions:read permission.""")
      # optional parameters
      parser.add_argument(
          """--targets""",
          default="""DeprecationWarning,UserWarning,FutureWarning""",
          type=list_str,
          help="""Comma-separated list of target warning(s) which we want to extract.""",
      )
      parser.add_argument(
          """--from_gh""",
          action="""store_true""",
          help="""If running from a GitHub action workflow and collecting warnings from its artifacts.""",
      )

      __magic_name__  :  Optional[Any]					=       parser.parse_args()

      __magic_name__  :  Optional[int]					=       args.from_gh
      if from_gh:
            # The artifacts have to be downloaded using `actions/download-artifact@v3`
            pass
      else:
            os.makedirs(args.output_dir, exist_ok=True)

            # get download links
            __magic_name__  :  Optional[Any]					=       get_artifacts_links(args.workflow_run_id, token=args.token)
            with open(os.path.join(args.output_dir, """artifacts.json"""), """w""", encoding="""UTF-8""") as fp:
                  json.dump(artifacts, fp, ensure_ascii=False, indent=4)

            # download artifacts
            for idx, (name, url) in enumerate(artifacts.items()):
                  print(name)
                  print(url)
                  print("""=""" * 80)
                  download_artifact(name, url, args.output_dir, args.token)
                  # Be gentle to GitHub
                  time.sleep(1)

    # extract warnings from artifacts
      __magic_name__  :  Optional[Any]					=       extract_warnings(args.output_dir, args.targets)
      __magic_name__  :  Any					=       sorted(selected_warnings)
      with open(os.path.join(args.output_dir, """selected_warnings.json"""), """w""", encoding="""UTF-8""") as fp:
            json.dump(selected_warnings, fp, ensure_ascii=False, indent=4)

