



'''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		):



 '''simple docstring'''







 def    a__    (self					,  A					,  A			)       ->	Optional[int]:


      """simple docstring"""



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



 def    a__    (self			)       ->	Tuple:


      """simple docstring"""



      _a           =    None
      _a           =    20

      _a           =    self._get_uniform_logits(batch_size=2					,  length=A			)

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

      _a           =    FlaxTemperatureLogitsWarper(temperature=0.5			)
      _a           =    FlaxTemperatureLogitsWarper(temperature=1.3			)

      _a           =    jax.nn.softmax(temp_dist_warper_sharper(A					,  scores.copy()					,  cur_len=A			)					,  axis=-1			)
      _a           =    jax.nn.softmax(temp_dist_warper_smoother(A					,  scores.copy()					,  cur_len=A			)					,  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]:


      """simple docstring"""



      _a           =    None
      _a           =    10
      _a           =    2

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

      _a           =    FlaxTopKLogitsWarper(3			)

      _a           =    top_k_warp(A					,  A					,  cur_len=A			)

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

      # 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			)       ->	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.15, 0.3, 0.3, 0.25]]			)			)

      _a           =    FlaxTopPLogitsWarper(0.8			)
      _a           =    np.exp(top_p_warp(A					,  A					,  cur_len=A			)			)

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

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

      # make ramp_logits more extreme
      _a           =    ramp_logits[1] * 100.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(A					,  A					,  cur_len=A			)

      # 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			)       ->	Tuple:


      """simple docstring"""



      _a           =    20
      _a           =    4
      _a           =    0

      _a           =    FlaxMinLengthLogitsProcessor(min_length=10					,  eos_token_id=A			)

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



 def    a__    (self			)       ->	Optional[int]:


      """simple docstring"""



      _a           =    20
      _a           =    4
      _a           =    0

      _a           =    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=A			)

      # 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(A					,  A			)
      _a           =    logits_processor(A					,  A					,  cur_len=A			)
      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(A					,  A			)
      _a           =    logits_processor(A					,  A					,  cur_len=A			)
      self.assertFalse(jnp.isinf(A			).any()			)



 def    a__    (self			)       ->	Optional[int]:


      """simple docstring"""



      _a           =    20
      _a           =    4
      _a           =    0
      _a           =    5

      _a           =    FlaxForcedEOSTokenLogitsProcessor(max_length=A					,  eos_token_id=A			)

      # 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(A					,  A			)
      _a           =    logits_processor(A					,  A					,  cur_len=A			)
      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(A					,  A			)
      _a           =    logits_processor(A					,  A					,  cur_len=A			)
      self.assertFalse(jnp.isinf(A			).any()			)



 def    a__    (self			)       ->	List[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)					,  A			)
      _a           =    input_ids.copy()

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

      _a           =    10

      # no processor list
      _a           =    temp_dist_warp(A					,  A					,  cur_len=A			)
      _a           =    top_k_warp(A					,  A					,  cur_len=A			)
      _a           =    top_p_warp(A					,  A					,  cur_len=A			)
      _a           =    min_dist_proc(A					,  A					,  cur_len=A			)
      _a           =    bos_dist_proc(A					,  A					,  cur_len=A			)
      _a           =    eos_dist_proc(A					,  A					,  cur_len=A			)

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

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

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



 def    a__    (self			)       ->	Optional[int]:


      """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)					,  A			)
      _a           =    input_ids.copy()

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

      _a           =    10

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

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

      _a           =    jax.jit(A			)
      _a           =    jax.jit(A			)

      _a           =    jitted_run_no_processor_list(A					,  A					,  A			)
      _a           =    jitted_run_processor_list(A					,  A					,  A			)

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

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










from __future__ import annotations

from scipy.special import comb  # type: ignore







class       __A							:




      """simple docstring"""







      def __init__(		self   ,				a__):



             """simple docstring"""





             _lowerCamelCase				:  str						=		list_of_points
             # Degree determines the flexibility of the curve.
             # Degree = 1 will produce a straight line.
             _lowerCamelCase				:  List[Any]						=		len(a__) - 1


      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             assert 0 <= t <= 1, "Time t must be between 0 and 1."
             _lowerCamelCase				:  list[float]						=		[]
             for i in range(len(self.list_of_points)):
                    # basis function for each i
                    output_values.append(
                        comb(self.degree   ,				a__) * ((1 - t) ** (self.degree - i)) * (t**i))
             # the basis must sum up to 1 for it to produce a valid Bezier curve.
             assert round(sum(a__)   ,				5) == 1
             return output_values


      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             assert 0 <= t <= 1, "Time t must be between 0 and 1."

             _lowerCamelCase				:  Any						=		self.basis_function(a__)
             _lowerCamelCase				:  List[Any]						=		0.0
             _lowerCamelCase				:  List[Any]						=		0.0
             for i in range(len(self.list_of_points)):
                    # For all points, sum up the product of i-th basis function and i-th point.
                    x += basis_function[i] * self.list_of_points[i][0]
                    y += basis_function[i] * self.list_of_points[i][1]
             return (x, y)



      def 					__snake_case      (		self   ,				a__ = 0.01):



             """simple docstring"""





             from matplotlib import pyplot as plt  # type: ignore

             _lowerCamelCase				:  list[float]						=		[]  # x coordinates of points to plot
             _lowerCamelCase				:  list[float]						=		[]  # y coordinates of points to plot

             _lowerCamelCase				:  List[Any]						=		0.0
             while t <= 1:
                    _lowerCamelCase				:  Tuple						=		self.bezier_curve_function(a__)
                    to_plot_x.append(value[0])
                    to_plot_y.append(value[1])
                    t += step_size

             _lowerCamelCase				:  List[Any]						=		[i[0] for i in self.list_of_points]
             _lowerCamelCase				:  Union[str, Any]						=		[i[1] for i in self.list_of_points]

             plt.plot(
                 a__   ,				a__   ,				color='''blue'''   ,				label='''Curve of Degree ''' + str(self.degree)   ,				)
             plt.scatter(a__   ,				a__   ,				color='''red'''   ,				label='''Control Points''')
             plt.legend()
             plt.show()


if __name__ == "__main__":
       import doctest

       doctest.testmod()

       BezierCurve([(1, 2), (3, 5)]).plot_curve()  # degree 1
       BezierCurve([(0, 0), (5, 5), (5, 0)]).plot_curve()  # degree 2
       BezierCurve([(0, 0), (5, 5), (5, 0), (2.5, -2.5)]).plot_curve()  # degree 3

