

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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =None
             __A    =2_0

             __A    =self._get_uniform_logits(batch_size=2 ,   length=lowercase__     )

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

             __A    =FlaxTemperatureLogitsWarper(temperature=0.5     )
             __A    =FlaxTemperatureLogitsWarper(temperature=1.3     )

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







             '''simple docstring'''



             __A    =None
             __A    =1_0
             __A    =2

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

             __A    =FlaxTopKLogitsWarper(3     )

             __A    =top_k_warp(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )

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

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







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

             __A    =FlaxTopPLogitsWarper(0.8     )
             __A    =np.exp(top_p_warp(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )     )

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

             # check edge cases with negative and extreme logits
             __A    =np.broadcast_to(np.arange(lowercase__     )[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(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )

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







             '''simple docstring'''



             __A    =2_0
             __A    =4
             __A    =0

             __A    =FlaxMinLengthLogitsProcessor(min_length=1_0 ,   eos_token_id=lowercase__     )

             # 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(lowercase__ ,   lowercase__     )
             __A    =min_dist_processor(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )
             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(lowercase__ ,   lowercase__     )
             __A    =1_5
             __A    =min_dist_processor(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )
             self.assertFalse(jnp.isinf(lowercase__     ).any()     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =2_0
             __A    =4
             __A    =0

             __A    =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase__     )

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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =2_0
             __A    =4
             __A    =0
             __A    =5

             __A    =FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase__ ,   eos_token_id=lowercase__     )

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

       def     __UpperCamelCase     (							self     ):







             '''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) ,   lowercase__     )
             __A    =input_ids.copy()

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

             __A    =1_0

             # no processor list
             __A    =temp_dist_warp(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )
             __A    =top_k_warp(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )
             __A    =top_p_warp(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )
             __A    =min_dist_proc(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )
             __A    =bos_dist_proc(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )
             __A    =eos_dist_proc(lowercase__ ,   lowercase__ ,   cur_len=lowercase__     )

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

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

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


       def     __UpperCamelCase     (							self     ):







             '''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) ,   lowercase__     )
             __A    =input_ids.copy()

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

             __A    =1_0

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

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

             __A    =jax.jit(lowercase__     )
             __A    =jax.jit(lowercase__     )

             __A    =jitted_run_no_processor_list(lowercase__ ,   lowercase__ ,   lowercase__     )
             __A    =jitted_run_processor_list(lowercase__ ,   lowercase__ ,   lowercase__     )

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

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





# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..utils import cached_file


# docstyle-ignore
_lowerCamelCase      :					Optional[int]					    =      '''
Human: <<task>>

Assistant: '''


_lowerCamelCase      :					str					    =      '''huggingface-tools/default-prompts'''
_lowerCamelCase      :					Optional[int]					    =      {'''chat''': '''chat_prompt_template.txt''', '''run''': '''run_prompt_template.txt'''}





def      A__					(			__A  :						Dict     ,				__A  :						Dict     ,				__A  :						Optional[Any]="run"			)					->Any:
      if prompt_or_repo_id is None:
            __A    =DEFAULT_PROMPTS_REPO

      # prompt is considered a repo ID when it does not contain any kind of space
      if re.search('''\\s'''     ,				__A			) is not None:
            return prompt_or_repo_id

      __A    =cached_file(
          __A     ,				PROMPT_FILES[mode]     ,				repo_type='''dataset'''     ,				user_agent={'''agent''': agent_name}			)
      with open(__A     ,				'''r'''     ,				encoding='''utf-8'''			) as f:
            return f.read()


