


'''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 __a						(					self     :							int				,							SCREAMING_SNAKE_CASE__     :							int				,							SCREAMING_SNAKE_CASE__     :							int						):





       '''simple docstring'''

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

  def __a						(					self     :							Dict						):





       '''simple docstring'''

       __a       =   None
       __a       =   2_0

       __a       =   self._get_uniform_logits(batch_size=2				,							length=SCREAMING_SNAKE_CASE__						)

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

       # compute softmax
       __a       =   jax.nn.softmax(SCREAMING_SNAKE_CASE__				,							axis=-1						)

       __a       =   FlaxTemperatureLogitsWarper(temperature=0.5						)
       __a       =   FlaxTemperatureLogitsWarper(temperature=1.3						)

       __a       =   jax.nn.softmax(temp_dist_warper_sharper(SCREAMING_SNAKE_CASE__				,							scores.copy()				,							cur_len=SCREAMING_SNAKE_CASE__						)				,							axis=-1						)
       __a       =   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 __a						(					self     :							Dict						):





       '''simple docstring'''

       __a       =   None
       __a       =   1_0
       __a       =   2

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

       __a       =   FlaxTopKLogitsWarper(3						)

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

       __a       =   np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE__						)[None, :]				,							(batch_size, length)						).copy()
       __a       =   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 __a						(					self     :							Any						):





       '''simple docstring'''

       __a       =   None
       __a       =   1_0
       __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(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
       __a       =   np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]						)
       self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							atol=1E-3						)						)

       # check edge cases with negative and extreme logits
       __a       =   np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE__						)[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(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 __a						(					self     :							Optional[Any]						):





       '''simple docstring'''

       __a       =   2_0
       __a       =   4
       __a       =   0

       __a       =   FlaxMinLengthLogitsProcessor(min_length=1_0				,							eos_token_id=SCREAMING_SNAKE_CASE__						)

       # check that min length is applied at length 5
       __a       =   ids_tensor((batch_size, 2_0)				,							vocab_size=2_0						)
       __a       =   5
       __a       =   self._get_uniform_logits(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						)
       __a       =   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
       __a       =   self._get_uniform_logits(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						)
       __a       =   1_5
       __a       =   min_dist_processor(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
       self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE__						).any()						)

  def __a						(					self     :							List[str]						):





       '''simple docstring'''

       __a       =   2_0
       __a       =   4
       __a       =   0

       __a       =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE__						)

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

  def __a						(					self     :							Optional[Any]						):





       '''simple docstring'''

       __a       =   2_0
       __a       =   4
       __a       =   0
       __a       =   5

       __a       =   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
       __a       =   ids_tensor((batch_size, 4)				,							vocab_size=2_0						)
       __a       =   4
       __a       =   self._get_uniform_logits(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						)
       __a       =   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
       __a       =   3
       __a       =   self._get_uniform_logits(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						)
       __a       =   logits_processor(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
       self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE__						).any()						)

  def __a						(					self     :							Tuple						):





       '''simple docstring'''

       __a       =   4
       __a       =   1_0
       __a       =   1_5
       __a       =   2
       __a       =   1
       __a       =   1_5

       # dummy input_ids and scores
       __a       =   ids_tensor((batch_size, sequence_length)				,							SCREAMING_SNAKE_CASE__						)
       __a       =   input_ids.copy()

       __a       =   self._get_uniform_logits(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						)
       __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=1_0				,							eos_token_id=SCREAMING_SNAKE_CASE__						)
       __a       =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE__						)
       __a       =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE__				,							eos_token_id=SCREAMING_SNAKE_CASE__						)

       __a       =   1_0

       # no processor list
       __a       =   temp_dist_warp(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
       __a       =   top_k_warp(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
       __a       =   top_p_warp(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
       __a       =   min_dist_proc(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
       __a       =   bos_dist_proc(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
       __a       =   eos_dist_proc(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)

       # 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(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 __a						(					self     :							int						):





       '''simple docstring'''

       __a       =   4
       __a       =   1_0
       __a       =   1_5
       __a       =   2
       __a       =   1
       __a       =   1_5

       # dummy input_ids and scores
       __a       =   ids_tensor((batch_size, sequence_length)				,							SCREAMING_SNAKE_CASE__						)
       __a       =   input_ids.copy()

       __a       =   self._get_uniform_logits(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						)
       __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=1_0				,							eos_token_id=SCREAMING_SNAKE_CASE__						)
       __a       =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE__						)
       __a       =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE__				,							eos_token_id=SCREAMING_SNAKE_CASE__						)

       __a       =   1_0

       # no processor list
       def run_no_processor_list(SCREAMING_SNAKE_CASE__     :							str				,							SCREAMING_SNAKE_CASE__     :							Dict				,							SCREAMING_SNAKE_CASE__     :							Optional[Any]						):
            __a       =   temp_dist_warp(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
            __a       =   top_k_warp(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
            __a       =   top_p_warp(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
            __a       =   min_dist_proc(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
            __a       =   bos_dist_proc(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
            __a       =   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__     :							List[str]				,							SCREAMING_SNAKE_CASE__     :							Optional[int]				,							SCREAMING_SNAKE_CASE__     :							Optional[Any]						):
            __a       =   FlaxLogitsProcessorList(
                [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
            __a       =   processor(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							cur_len=SCREAMING_SNAKE_CASE__						)
            return scores

       __a       =   jax.jit(SCREAMING_SNAKE_CASE__						)
       __a       =   jax.jit(SCREAMING_SNAKE_CASE__						)

       __a       =   jitted_run_no_processor_list(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						)
       __a       =   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'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


SCREAMING_SNAKE_CASE_      =			{
    'configuration_bridgetower': [
        'BRIDGETOWER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'BridgeTowerConfig',
        'BridgeTowerTextConfig',
        'BridgeTowerVisionConfig',
    ],
    'processing_bridgetower': ['BridgeTowerProcessor'],
}

try:
 if not is_vision_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE_      =			['BridgeTowerImageProcessor']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE_      =			[
     'BRIDGETOWER_PRETRAINED_MODEL_ARCHIVE_LIST',
     'BridgeTowerForContrastiveLearning',
     'BridgeTowerForImageAndTextRetrieval',
     'BridgeTowerForMaskedLM',
     'BridgeTowerModel',
     'BridgeTowerPreTrainedModel',
 ]


if TYPE_CHECKING:
 from .configuration_bridgetower import (
     BRIDGETOWER_PRETRAINED_CONFIG_ARCHIVE_MAP,
     BridgeTowerConfig,
     BridgeTowerTextConfig,
     BridgeTowerVisionConfig,
 )
 from .processing_bridgetower import BridgeTowerProcessor

 try:
  if not is_vision_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .image_processing_bridgetower import BridgeTowerImageProcessor

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_bridgetower import (
      BRIDGETOWER_PRETRAINED_MODEL_ARCHIVE_LIST,
      BridgeTowerForContrastiveLearning,
      BridgeTowerForImageAndTextRetrieval,
      BridgeTowerForMaskedLM,
      BridgeTowerModel,
      BridgeTowerPreTrainedModel,
  )


else:
 import sys

 SCREAMING_SNAKE_CASE_      =			_LazyModule(__name__, globals()['__file__'], _import_structure)






