


'''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   __UpperCamelCase   (    self		,      lowerCamelCase_		,      lowerCamelCase_      )      ->    List[Any]:
    _a					:  str						=  jnp.ones((batch_size, length)      ) / length
    return scores




 def   __UpperCamelCase   (    self      )      ->    int:
    _a					:  str						=  None
    _a					:  List[str]						=  2_0

    _a					:  List[str]						=  self._get_uniform_logits(batch_size=2		,      length=lowerCamelCase_      )

    # tweak scores to not be uniform anymore
    _a					:  Union[str, Any]						=  scores.at[1, 5].set((1 / length) + 0.1      )  # peak, 1st batch
    _a					:  Optional[int]						=  scores.at[1, 1_0].set((1 / length) - 0.4      )  # valley, 1st batch

    # compute softmax
    _a					:  List[str]						=  jax.nn.softmax(lowerCamelCase_		,      axis=-1      )

    _a					:  Optional[Any]						=  FlaxTemperatureLogitsWarper(temperature=0.5      )
    _a					:  Any						=  FlaxTemperatureLogitsWarper(temperature=1.3      )

    _a					:  Optional[Any]						=  jax.nn.softmax(temp_dist_warper_sharper(lowerCamelCase_		,      scores.copy()		,      cur_len=lowerCamelCase_      )		,      axis=-1      )
    _a					:  Tuple						=  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   __UpperCamelCase   (    self      )      ->    int:
    _a					:  List[Any]						=  None
    _a					:  Any						=  1_0
    _a					:  Tuple						=  2

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

    _a					:  int						=  FlaxTopKLogitsWarper(3      )

    _a					:  int						=  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					:  Optional[int]						=  5
    _a					:  str						=  FlaxTopKLogitsWarper(top_k=1		,      filter_value=0.0		,      min_tokens_to_keep=3      )

    _a					:  str						=  np.broadcast_to(np.arange(lowerCamelCase_      )[None, :]		,      (batch_size, length)      ).copy()
    _a					:  Union[str, Any]						=  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   __UpperCamelCase   (    self      )      ->    Union[str, Any]:
    _a					:  Optional[Any]						=  None
    _a					:  str						=  1_0
    _a					:  Any						=  2

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

    _a					:  List[str]						=  FlaxTopPLogitsWarper(0.8      )
    _a					:  Any						=  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					:  Dict						=  np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]      )
    self.assertTrue(np.allclose(lowerCamelCase_		,      lowerCamelCase_		,      atol=1e-3      )      )

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

    # make ramp_logits more extreme
    _a					:  List[Any]						=  ramp_logits[1] * 100.0

    # make sure at least 2 tokens are kept
    _a					:  Union[str, Any]						=  FlaxTopPLogitsWarper(0.9		,      min_tokens_to_keep=2		,      filter_value=0.0      )
    _a					:  Any						=  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   __UpperCamelCase   (    self      )      ->    List[str]:
    _a					:  Union[str, Any]						=  2_0
    _a					:  List[str]						=  4
    _a					:  str						=  0

    _a					:  List[str]						=  FlaxMinLengthLogitsProcessor(min_length=1_0		,      eos_token_id=lowerCamelCase_      )

    # check that min length is applied at length 5
    _a					:  List[str]						=  ids_tensor((batch_size, 2_0)		,      vocab_size=2_0      )
    _a					:  str						=  5
    _a					:  List[str]						=  self._get_uniform_logits(lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  List[str]						=  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					:  str						=  self._get_uniform_logits(lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  List[Any]						=  1_5
    _a					:  List[Any]						=  min_dist_processor(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
    self.assertFalse(jnp.isinf(lowerCamelCase_      ).any()      )




 def   __UpperCamelCase   (    self      )      ->    Optional[int]:
    _a					:  Tuple						=  2_0
    _a					:  int						=  4
    _a					:  int						=  0

    _a					:  int						=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_      )

    # check that all scores are -inf except the bos_token_id score
    _a					:  Tuple						=  ids_tensor((batch_size, 1)		,      vocab_size=2_0      )
    _a					:  Optional[Any]						=  1
    _a					:  int						=  self._get_uniform_logits(lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  Tuple						=  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					:  Any						=  3
    _a					:  List[Any]						=  self._get_uniform_logits(lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  int						=  logits_processor(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
    self.assertFalse(jnp.isinf(lowerCamelCase_      ).any()      )




 def   __UpperCamelCase   (    self      )      ->    List[Any]:
    _a					:  List[Any]						=  2_0
    _a					:  List[Any]						=  4
    _a					:  Union[str, Any]						=  0
    _a					:  Any						=  5

    _a					:  Optional[int]						=  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					:  List[str]						=  ids_tensor((batch_size, 4)		,      vocab_size=2_0      )
    _a					:  Union[str, Any]						=  4
    _a					:  str						=  self._get_uniform_logits(lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  Union[str, Any]						=  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					:  str						=  3
    _a					:  Any						=  self._get_uniform_logits(lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  List[str]						=  logits_processor(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
    self.assertFalse(jnp.isinf(lowerCamelCase_      ).any()      )




 def   __UpperCamelCase   (    self      )      ->    List[Any]:
    _a					:  Dict						=  4
    _a					:  Any						=  1_0
    _a					:  Optional[int]						=  1_5
    _a					:  List[str]						=  2
    _a					:  List[str]						=  1
    _a					:  Union[str, Any]						=  1_5

    # dummy input_ids and scores
    _a					:  int						=  ids_tensor((batch_size, sequence_length)		,      lowerCamelCase_      )
    _a					:  Union[str, Any]						=  input_ids.copy()

    _a					:  Union[str, Any]						=  self._get_uniform_logits(lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  List[str]						=  scores.copy()

    # instantiate all dist processors
    _a					:  str						=  FlaxTemperatureLogitsWarper(temperature=0.5      )
    _a					:  Optional[Any]						=  FlaxTopKLogitsWarper(3      )
    _a					:  Optional[Any]						=  FlaxTopPLogitsWarper(0.8      )

    # instantiate all logits processors
    _a					:  int						=  FlaxMinLengthLogitsProcessor(min_length=1_0		,      eos_token_id=lowerCamelCase_      )
    _a					:  Optional[Any]						=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_      )
    _a					:  Any						=  FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase_		,      eos_token_id=lowerCamelCase_      )

    _a					:  int						=  1_0

    # no processor list
    _a					:  Dict						=  temp_dist_warp(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
    _a					:  str						=  top_k_warp(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
    _a					:  List[str]						=  top_p_warp(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
    _a					:  List[str]						=  min_dist_proc(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
    _a					:  Optional[Any]						=  bos_dist_proc(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
    _a					:  List[str]						=  eos_dist_proc(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )

    # with processor list
    _a					:  Optional[Any]						=  FlaxLogitsProcessorList(
        [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]      )
    _a					:  Optional[int]						=  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   __UpperCamelCase   (    self      )      ->    Union[str, Any]:
    _a					:  Any						=  4
    _a					:  str						=  1_0
    _a					:  int						=  1_5
    _a					:  Optional[int]						=  2
    _a					:  Optional[int]						=  1
    _a					:  Union[str, Any]						=  1_5

    # dummy input_ids and scores
    _a					:  int						=  ids_tensor((batch_size, sequence_length)		,      lowerCamelCase_      )
    _a					:  Union[str, Any]						=  input_ids.copy()

    _a					:  List[Any]						=  self._get_uniform_logits(lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  Optional[int]						=  scores.copy()

    # instantiate all dist processors
    _a					:  Optional[Any]						=  FlaxTemperatureLogitsWarper(temperature=0.5      )
    _a					:  str						=  FlaxTopKLogitsWarper(3      )
    _a					:  Optional[Any]						=  FlaxTopPLogitsWarper(0.8      )

    # instantiate all logits processors
    _a					:  int						=  FlaxMinLengthLogitsProcessor(min_length=1_0		,      eos_token_id=lowerCamelCase_      )
    _a					:  Optional[Any]						=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_      )
    _a					:  Optional[int]						=  FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase_		,      eos_token_id=lowerCamelCase_      )

    _a					:  int						=  1_0

    # no processor list
    def run_no_processor_list(lowerCamelCase_		,      lowerCamelCase_		,      lowerCamelCase_      ):
       _a					:  Tuple						=  temp_dist_warp(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
       _a					:  Optional[int]						=  top_k_warp(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
       _a					:  int						=  top_p_warp(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
       _a					:  Any						=  min_dist_proc(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
       _a					:  Any						=  bos_dist_proc(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
       _a					:  int						=  eos_dist_proc(lowerCamelCase_		,      lowerCamelCase_		,      cur_len=lowerCamelCase_      )
       return scores

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

    _a					:  Union[str, Any]						=  jax.jit(lowerCamelCase_      )
    _a					:  int						=  jax.jit(lowerCamelCase_      )

    _a					:  str						=  jitted_run_no_processor_list(lowerCamelCase_		,      lowerCamelCase_		,      lowerCamelCase_      )
    _a					:  Optional[Any]						=  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()      )






'''simple docstring'''




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


UpperCAmelCase_      :		Dict          =		logging.get_logger(__name__)

UpperCAmelCase_      :		List[str]          =		{
    "uclanlp/visualbert-vqa": "https://huggingface.co/uclanlp/visualbert-vqa/resolve/main/config.json",
    "uclanlp/visualbert-vqa-pre": "https://huggingface.co/uclanlp/visualbert-vqa-pre/resolve/main/config.json",
    "uclanlp/visualbert-vqa-coco-pre": (
        "https://huggingface.co/uclanlp/visualbert-vqa-coco-pre/resolve/main/config.json"
    ),
    "uclanlp/visualbert-vcr": "https://huggingface.co/uclanlp/visualbert-vcr/resolve/main/config.json",
    "uclanlp/visualbert-vcr-pre": "https://huggingface.co/uclanlp/visualbert-vcr-pre/resolve/main/config.json",
    "uclanlp/visualbert-vcr-coco-pre": (
        "https://huggingface.co/uclanlp/visualbert-vcr-coco-pre/resolve/main/config.json"
    ),
    "uclanlp/visualbert-nlvr2": "https://huggingface.co/uclanlp/visualbert-nlvr2/resolve/main/config.json",
    "uclanlp/visualbert-nlvr2-pre": "https://huggingface.co/uclanlp/visualbert-nlvr2-pre/resolve/main/config.json",
    "uclanlp/visualbert-nlvr2-coco-pre": (
        "https://huggingface.co/uclanlp/visualbert-nlvr2-coco-pre/resolve/main/config.json"
    )
    # See all VisualBERT models at https://huggingface.co/models?filter=visual_bert
}







class 					a			(				snake_case__  ):
 '''simple docstring'''




 __lowerCAmelCase						: int							  =		"""visual_bert"""




 def __init__(    self		,      lowerCamelCase_=3_0_5_2_2		,      lowerCamelCase_=7_6_8		,      lowerCamelCase_=5_1_2		,      lowerCamelCase_=1_2		,      lowerCamelCase_=1_2		,      lowerCamelCase_=3_0_7_2		,      lowerCamelCase_="gelu"		,      lowerCamelCase_=0.1		,      lowerCamelCase_=0.1		,      lowerCamelCase_=5_1_2		,      lowerCamelCase_=2		,      lowerCamelCase_=0.02		,      lowerCamelCase_=1e-12		,      lowerCamelCase_=False		,      lowerCamelCase_=True		,      lowerCamelCase_=1		,      lowerCamelCase_=0		,      lowerCamelCase_=2		,      **lowerCamelCase_		,      )      ->    Optional[Any]:
    super().__init__(pad_token_id=lowerCamelCase_		,      bos_token_id=lowerCamelCase_		,      eos_token_id=lowerCamelCase_		,      **lowerCamelCase_      )

    _a					:  List[str]						=  vocab_size
    _a					:  Union[str, Any]						=  max_position_embeddings
    _a					:  Tuple						=  hidden_size
    _a					:  Tuple						=  visual_embedding_dim
    _a					:  Optional[Any]						=  num_hidden_layers
    _a					:  str						=  num_attention_heads
    _a					:  Union[str, Any]						=  intermediate_size
    _a					:  List[str]						=  hidden_act
    _a					:  str						=  hidden_dropout_prob
    _a					:  List[Any]						=  attention_probs_dropout_prob
    _a					:  Optional[Any]						=  initializer_range
    _a					:  List[Any]						=  type_vocab_size
    _a					:  Any						=  layer_norm_eps
    _a					:  Tuple						=  bypass_transformer
    _a					:  str						=  special_visual_initialize


