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




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       Any  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       int  )				->Optional[Any]:
        UpperCAmelCase_            =  jnp.ones((batch_size, length)  ) / length
        return scores






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Tuple:
        UpperCAmelCase_            =  None
        UpperCAmelCase_            =  20

        UpperCAmelCase_            =  self._get_uniform_logits(batch_size=2  ,							length=UpperCAmelCase__  )

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

        UpperCAmelCase_            =  FlaxTemperatureLogitsWarper(temperature=0.5  )
        UpperCAmelCase_            =  FlaxTemperatureLogitsWarper(temperature=1.3  )

        UpperCAmelCase_            =  jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase__  ,							scores.copy()  ,							cur_len=UpperCAmelCase__  )  ,							axis=-1  )
        UpperCAmelCase_            =  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  lowerCAmelCase__ (      self			:       List[str]  )				->str:
        UpperCAmelCase_            =  None
        UpperCAmelCase_            =  10
        UpperCAmelCase_            =  2

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

        UpperCAmelCase_            =  FlaxTopKLogitsWarper(3  )

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

        UpperCAmelCase_            =  np.broadcast_to(np.arange(UpperCAmelCase__  )[None, :]  ,							(batch_size, length)  ).copy()
        UpperCAmelCase_            =  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  lowerCAmelCase__ (      self			:       Any  )				->Union[str, Any]:
        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(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )  )

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

        # check edge cases with negative and extreme logits
        UpperCAmelCase_            =  np.broadcast_to(np.arange(UpperCAmelCase__  )[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(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  lowerCAmelCase__ (      self			:       Dict  )				->List[Any]:
        UpperCAmelCase_            =  20
        UpperCAmelCase_            =  4
        UpperCAmelCase_            =  0

        UpperCAmelCase_            =  FlaxMinLengthLogitsProcessor(min_length=10  ,							eos_token_id=UpperCAmelCase__  )

        # 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(UpperCAmelCase__  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  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
        UpperCAmelCase_            =  self._get_uniform_logits(UpperCAmelCase__  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  15
        UpperCAmelCase_            =  min_dist_processor(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
        self.assertFalse(jnp.isinf(UpperCAmelCase__  ).any()  )






  def  lowerCAmelCase__ (      self			:       Any  )				->str:
        UpperCAmelCase_            =  20
        UpperCAmelCase_            =  4
        UpperCAmelCase_            =  0

        UpperCAmelCase_            =  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__  )

        # 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(UpperCAmelCase__  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  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
        UpperCAmelCase_            =  3
        UpperCAmelCase_            =  self._get_uniform_logits(UpperCAmelCase__  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  logits_processor(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
        self.assertFalse(jnp.isinf(UpperCAmelCase__  ).any()  )






  def  lowerCAmelCase__ (      self			:       Tuple  )				->Dict:
        UpperCAmelCase_            =  20
        UpperCAmelCase_            =  4
        UpperCAmelCase_            =  0
        UpperCAmelCase_            =  5

        UpperCAmelCase_            =  FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__  ,							eos_token_id=UpperCAmelCase__  )

        # 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(UpperCAmelCase__  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  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
        UpperCAmelCase_            =  3
        UpperCAmelCase_            =  self._get_uniform_logits(UpperCAmelCase__  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  logits_processor(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
        self.assertFalse(jnp.isinf(UpperCAmelCase__  ).any()  )






  def  lowerCAmelCase__ (      self			:       Tuple  )				->Optional[int]:
        UpperCAmelCase_            =  4
        UpperCAmelCase_            =  10
        UpperCAmelCase_            =  15
        UpperCAmelCase_            =  2
        UpperCAmelCase_            =  1
        UpperCAmelCase_            =  15

        # dummy input_ids and scores
        UpperCAmelCase_            =  ids_tensor((batch_size, sequence_length)  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  input_ids.copy()

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

        UpperCAmelCase_            =  10

        # no processor list
        UpperCAmelCase_            =  temp_dist_warp(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
        UpperCAmelCase_            =  top_k_warp(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
        UpperCAmelCase_            =  top_p_warp(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
        UpperCAmelCase_            =  min_dist_proc(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
        UpperCAmelCase_            =  bos_dist_proc(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
        UpperCAmelCase_            =  eos_dist_proc(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )

        # 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(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  lowerCAmelCase__ (      self			:       List[str]  )				->Optional[int]:
        UpperCAmelCase_            =  4
        UpperCAmelCase_            =  10
        UpperCAmelCase_            =  15
        UpperCAmelCase_            =  2
        UpperCAmelCase_            =  1
        UpperCAmelCase_            =  15

        # dummy input_ids and scores
        UpperCAmelCase_            =  ids_tensor((batch_size, sequence_length)  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  input_ids.copy()

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

        UpperCAmelCase_            =  10

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

        # with processor list
        def run_processor_list(UpperCAmelCase__			:       Any  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       Tuple  ):
              UpperCAmelCase_            =  FlaxLogitsProcessorList(
                  [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
              UpperCAmelCase_            =  processor(UpperCAmelCase__  ,							UpperCAmelCase__  ,							cur_len=UpperCAmelCase__  )
              return scores

        UpperCAmelCase_            =  jax.jit(UpperCAmelCase__  )
        UpperCAmelCase_            =  jax.jit(UpperCAmelCase__  )

        UpperCAmelCase_            =  jitted_run_no_processor_list(UpperCAmelCase__  ,							UpperCAmelCase__  ,							UpperCAmelCase__  )
        UpperCAmelCase_            =  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()  )






'''simple docstring'''
from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase__  :      int        =					logging.get_logger(__name__)

lowercase__  :      List[str]        =					{
    "huggingface/time-series-transformer-tourism-monthly": (
        "https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}







class lowerCamelCase							(	lowerCamelCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	'''time_series_transformer'''
  lowerCAmelCase__										=	{
      '''hidden_size''': '''d_model''',
      '''num_attention_heads''': '''encoder_attention_heads''',
      '''num_hidden_layers''': '''encoder_layers''',
  }






  def __init__(      self			:       Optional[int]  ,							UpperCAmelCase__			:       Optional[int] = None  ,							UpperCAmelCase__			:       Optional[int] = None  ,							UpperCAmelCase__			:       str = "student_t"  ,							UpperCAmelCase__			:       str = "nll"  ,							UpperCAmelCase__			:       int = 1  ,							UpperCAmelCase__			:       List[int] = [1, 2, 3, 4, 5, 6, 7]  ,							UpperCAmelCase__			:       Optional[Union[str, bool]] = "mean"  ,							UpperCAmelCase__			:       int = 0  ,							UpperCAmelCase__			:       int = 0  ,							UpperCAmelCase__			:       int = 0  ,							UpperCAmelCase__			:       int = 0  ,							UpperCAmelCase__			:       Optional[List[int]] = None  ,							UpperCAmelCase__			:       Optional[List[int]] = None  ,							UpperCAmelCase__			:       int = 32  ,							UpperCAmelCase__			:       int = 32  ,							UpperCAmelCase__			:       int = 2  ,							UpperCAmelCase__			:       int = 2  ,							UpperCAmelCase__			:       int = 2  ,							UpperCAmelCase__			:       int = 2  ,							UpperCAmelCase__			:       bool = True  ,							UpperCAmelCase__			:       str = "gelu"  ,							UpperCAmelCase__			:       int = 64  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       int = 100  ,							UpperCAmelCase__			:       float = 0.02  ,							UpperCAmelCase__			:       Union[str, Any]=True  ,							**UpperCAmelCase__			:       Tuple  ,							)				->Optional[int]:
        # time series specific configuration
        UpperCAmelCase_            =  prediction_length
        UpperCAmelCase_            =  context_length or prediction_length
        UpperCAmelCase_            =  distribution_output
        UpperCAmelCase_            =  loss
        UpperCAmelCase_            =  input_size
        UpperCAmelCase_            =  num_time_features
        UpperCAmelCase_            =  lags_sequence
        UpperCAmelCase_            =  scaling
        UpperCAmelCase_            =  num_dynamic_real_features
        UpperCAmelCase_            =  num_static_real_features
        UpperCAmelCase_            =  num_static_categorical_features
        if cardinality and num_static_categorical_features > 0:
              if len(UpperCAmelCase__  ) != num_static_categorical_features:
                    raise ValueError(
                        '''The cardinality should be a list of the same length as `num_static_categorical_features`'''  )
              UpperCAmelCase_            =  cardinality
        else:
              UpperCAmelCase_            =  [0]
        if embedding_dimension and num_static_categorical_features > 0:
              if len(UpperCAmelCase__  ) != num_static_categorical_features:
                    raise ValueError(
                        '''The embedding dimension should be a list of the same length as `num_static_categorical_features`'''  )
              UpperCAmelCase_            =  embedding_dimension
        else:
              UpperCAmelCase_            =  [min(50  ,							(cat + 1) // 2  ) for cat in self.cardinality]
        UpperCAmelCase_            =  num_parallel_samples

        # Transformer architecture configuration
        UpperCAmelCase_            =  input_size * len(UpperCAmelCase__  ) + self._number_of_features
        UpperCAmelCase_            =  d_model
        UpperCAmelCase_            =  encoder_attention_heads
        UpperCAmelCase_            =  decoder_attention_heads
        UpperCAmelCase_            =  encoder_ffn_dim
        UpperCAmelCase_            =  decoder_ffn_dim
        UpperCAmelCase_            =  encoder_layers
        UpperCAmelCase_            =  decoder_layers

        UpperCAmelCase_            =  dropout
        UpperCAmelCase_            =  attention_dropout
        UpperCAmelCase_            =  activation_dropout
        UpperCAmelCase_            =  encoder_layerdrop
        UpperCAmelCase_            =  decoder_layerdrop

        UpperCAmelCase_            =  activation_function
        UpperCAmelCase_            =  init_std

        UpperCAmelCase_            =  use_cache

        super().__init__(is_encoder_decoder=UpperCAmelCase__  ,							**UpperCAmelCase__  )






  @property
  def  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        return (
            sum(self.embedding_dimension  )
            + self.num_dynamic_real_features
            + self.num_time_features
            + self.num_static_real_features
            + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
        )





