




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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A			):




											"""simple docstring"""



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






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       None
											__lowerCAmelCase				   =       2_0

											__lowerCAmelCase				   =       self._get_uniform_logits(batch_size=2		,			length=_A			)

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

											# compute softmax
											__lowerCAmelCase				   =       jax.nn.softmax(_A		,			axis=-1			)

											__lowerCAmelCase				   =       FlaxTemperatureLogitsWarper(temperature=0.5			)
											__lowerCAmelCase				   =       FlaxTemperatureLogitsWarper(temperature=1.3			)

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




											"""simple docstring"""



											__lowerCAmelCase				   =       None
											__lowerCAmelCase				   =       1_0
											__lowerCAmelCase				   =       2

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

											__lowerCAmelCase				   =       FlaxTopKLogitsWarper(3			)

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

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




											"""simple docstring"""



											__lowerCAmelCase				   =       None
											__lowerCAmelCase				   =       1_0
											__lowerCAmelCase				   =       2

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

											__lowerCAmelCase				   =       FlaxTopPLogitsWarper(0.8			)
											__lowerCAmelCase				   =       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
											__lowerCAmelCase				   =       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
											__lowerCAmelCase				   =       np.broadcast_to(np.arange(_A			)[None, :]		,			(batch_size, vocab_size)			).copy() - (
											    vocab_size // 2
											)

											# make ramp_logits more extreme
											__lowerCAmelCase				   =       ramp_logits[1] * 1_00.0

											# make sure at least 2 tokens are kept
											__lowerCAmelCase				   =       FlaxTopPLogitsWarper(0.9		,			min_tokens_to_keep=2		,			filter_value=0.0			)
											__lowerCAmelCase				   =       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 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       2_0
											__lowerCAmelCase				   =       4
											__lowerCAmelCase				   =       0

											__lowerCAmelCase				   =       FlaxMinLengthLogitsProcessor(min_length=1_0		,			eos_token_id=_A			)

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






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       2_0
											__lowerCAmelCase				   =       4
											__lowerCAmelCase				   =       0

											__lowerCAmelCase				   =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A			)

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






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       2_0
											__lowerCAmelCase				   =       4
											__lowerCAmelCase				   =       0
											__lowerCAmelCase				   =       5

											__lowerCAmelCase				   =       FlaxForcedEOSTokenLogitsProcessor(max_length=_A		,			eos_token_id=_A			)

											# check that all scores are -inf except the eos_token_id when max_length is reached
											__lowerCAmelCase				   =       ids_tensor((batch_size, 4)		,			vocab_size=2_0			)
											__lowerCAmelCase				   =       4
											__lowerCAmelCase				   =       self._get_uniform_logits(_A		,			_A			)
											__lowerCAmelCase				   =       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
											__lowerCAmelCase				   =       3
											__lowerCAmelCase				   =       self._get_uniform_logits(_A		,			_A			)
											__lowerCAmelCase				   =       logits_processor(_A		,			_A		,			cur_len=_A			)
											self.assertFalse(jnp.isinf(_A			).any()			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       4
											__lowerCAmelCase				   =       1_0
											__lowerCAmelCase				   =       1_5
											__lowerCAmelCase				   =       2
											__lowerCAmelCase				   =       1
											__lowerCAmelCase				   =       1_5

											# dummy input_ids and scores
											__lowerCAmelCase				   =       ids_tensor((batch_size, sequence_length)		,			_A			)
											__lowerCAmelCase				   =       input_ids.copy()

											__lowerCAmelCase				   =       self._get_uniform_logits(_A		,			_A			)
											__lowerCAmelCase				   =       scores.copy()

											# instantiate all dist processors
											__lowerCAmelCase				   =       FlaxTemperatureLogitsWarper(temperature=0.5			)
											__lowerCAmelCase				   =       FlaxTopKLogitsWarper(3			)
											__lowerCAmelCase				   =       FlaxTopPLogitsWarper(0.8			)

											# instantiate all logits processors
											__lowerCAmelCase				   =       FlaxMinLengthLogitsProcessor(min_length=1_0		,			eos_token_id=_A			)
											__lowerCAmelCase				   =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A			)
											__lowerCAmelCase				   =       FlaxForcedEOSTokenLogitsProcessor(max_length=_A		,			eos_token_id=_A			)

											__lowerCAmelCase				   =       1_0

											# no processor list
											__lowerCAmelCase				   =       temp_dist_warp(_A		,			_A		,			cur_len=_A			)
											__lowerCAmelCase				   =       top_k_warp(_A		,			_A		,			cur_len=_A			)
											__lowerCAmelCase				   =       top_p_warp(_A		,			_A		,			cur_len=_A			)
											__lowerCAmelCase				   =       min_dist_proc(_A		,			_A		,			cur_len=_A			)
											__lowerCAmelCase				   =       bos_dist_proc(_A		,			_A		,			cur_len=_A			)
											__lowerCAmelCase				   =       eos_dist_proc(_A		,			_A		,			cur_len=_A			)

											# with processor list
											__lowerCAmelCase				   =       FlaxLogitsProcessorList(
											    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
											__lowerCAmelCase				   =       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 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       4
											__lowerCAmelCase				   =       1_0
											__lowerCAmelCase				   =       1_5
											__lowerCAmelCase				   =       2
											__lowerCAmelCase				   =       1
											__lowerCAmelCase				   =       1_5

											# dummy input_ids and scores
											__lowerCAmelCase				   =       ids_tensor((batch_size, sequence_length)		,			_A			)
											__lowerCAmelCase				   =       input_ids.copy()

											__lowerCAmelCase				   =       self._get_uniform_logits(_A		,			_A			)
											__lowerCAmelCase				   =       scores.copy()

											# instantiate all dist processors
											__lowerCAmelCase				   =       FlaxTemperatureLogitsWarper(temperature=0.5			)
											__lowerCAmelCase				   =       FlaxTopKLogitsWarper(3			)
											__lowerCAmelCase				   =       FlaxTopPLogitsWarper(0.8			)

											# instantiate all logits processors
											__lowerCAmelCase				   =       FlaxMinLengthLogitsProcessor(min_length=1_0		,			eos_token_id=_A			)
											__lowerCAmelCase				   =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A			)
											__lowerCAmelCase				   =       FlaxForcedEOSTokenLogitsProcessor(max_length=_A		,			eos_token_id=_A			)

											__lowerCAmelCase				   =       1_0

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

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

											__lowerCAmelCase				   =       jax.jit(_A			)
											__lowerCAmelCase				   =       jax.jit(_A			)

											__lowerCAmelCase				   =       jitted_run_no_processor_list(_A		,			_A		,			_A			)
											__lowerCAmelCase				   =       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()			)










import argparse
import os
import re

import packaging.version


UpperCamelCase__    			=			"""examples/"""
UpperCamelCase__    			=			{
    """examples""": (re.compile(R"""^check_min_version\(\"[^\"]+\"\)\s*$""", re.MULTILINE), """check_min_version(\"VERSION\")\n"""),
    """init""": (re.compile(R"""^__version__\s+=\s+\"([^\"]+)\"\s*$""", re.MULTILINE), """__version__ = \"VERSION\"\n"""),
    """setup""": (re.compile(R"""^(\s*)version\s*=\s*\"[^\"]+\",""", re.MULTILINE), R"""\1version=\"VERSION\","""),
    """doc""": (re.compile(R"""^(\s*)release\s*=\s*\"[^\"]+\"$""", re.MULTILINE), """release = \"VERSION\"\n"""),
}
UpperCamelCase__    			=			{
    """init""": """src/transformers/__init__.py""",
    """setup""": """setup.py""",
}
UpperCamelCase__    			=			"""README.md"""







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Optional[int]      ,      SCREAMING_SNAKE_CASE_							:				Any      ,      SCREAMING_SNAKE_CASE_							:				List[str]     ):
					with open(SCREAMING_SNAKE_CASE_      ,      "r"      ,      encoding="utf-8"      ,      newline="\n"     ) as f:
										__lowerCAmelCase				   =       f.read()
					__lowerCAmelCase   ,							__lowerCAmelCase				   =       REPLACE_PATTERNS[pattern]
					__lowerCAmelCase				   =       replace.replace("VERSION"      ,      SCREAMING_SNAKE_CASE_     )
					__lowerCAmelCase				   =       re_pattern.sub(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
					with open(SCREAMING_SNAKE_CASE_      ,      "w"      ,      encoding="utf-8"      ,      newline="\n"     ) as f:
										f.write(SCREAMING_SNAKE_CASE_     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				List[Any]     ):
					for folder, directories, fnames in os.walk(SCREAMING_SNAKE_CASE_     ):
										# Removing some of the folders with non-actively maintained examples from the walk
										if "research_projects" in directories:
															directories.remove("research_projects"     )
										if "legacy" in directories:
															directories.remove("legacy"     )
										for fname in fnames:
															if fname.endswith(".py"     ):
																				update_version_in_file(os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )      ,      SCREAMING_SNAKE_CASE_      ,      pattern="examples"     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Tuple      ,      SCREAMING_SNAKE_CASE_							:				Optional[int]=False     ):
					for pattern, fname in REPLACE_FILES.items():
										update_version_in_file(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
					if not patch:
										update_version_in_examples(SCREAMING_SNAKE_CASE_     )







def 	_a      (						):
					__lowerCAmelCase				   =       "🤗 Transformers currently provides the following architectures"
					__lowerCAmelCase				   =       "1. Want to contribute a new model?"
					with open(SCREAMING_SNAKE_CASE_      ,      "r"      ,      encoding="utf-8"      ,      newline="\n"     ) as f:
										__lowerCAmelCase				   =       f.readlines()

					# Find the start of the list.
					__lowerCAmelCase				   =       0
					while not lines[start_index].startswith(_start_prompt     ):
										start_index += 1
					start_index += 1

					__lowerCAmelCase				   =       start_index
					# Update the lines in the model list.
					while not lines[index].startswith(_end_prompt     ):
										if lines[index].startswith("1."     ):
															__lowerCAmelCase				   =       lines[index].replace(
															    "https://huggingface.co/docs/transformers/main/model_doc"      ,      "https://huggingface.co/docs/transformers/model_doc"      ,      )
										index += 1

					with open(SCREAMING_SNAKE_CASE_      ,      "w"      ,      encoding="utf-8"      ,      newline="\n"     ) as f:
										f.writelines(SCREAMING_SNAKE_CASE_     )







def 	_a      (						):
					with open(REPLACE_FILES["init"]      ,      "r"     ) as f:
										__lowerCAmelCase				   =       f.read()
					__lowerCAmelCase				   =       REPLACE_PATTERNS["init"][0].search(SCREAMING_SNAKE_CASE_     ).groups()[0]
					return packaging.version.parse(SCREAMING_SNAKE_CASE_     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				List[Any]=False     ):
					__lowerCAmelCase				   =       get_version()
					if patch and default_version.is_devrelease:
										raise ValueError("Can't create a patch version from the dev branch, checkout a released version!"     )
					if default_version.is_devrelease:
										__lowerCAmelCase				   =       default_version.base_version
					elif patch:
										__lowerCAmelCase				   =       F"""{default_version.major}.{default_version.minor}.{default_version.micro + 1}"""
					else:
										__lowerCAmelCase				   =       F"""{default_version.major}.{default_version.minor + 1}.0"""

					# Now let's ask nicely if that's the right one.
					__lowerCAmelCase				   =       input(F"""Which version are you releasing? [{default_version}]"""     )
					if len(SCREAMING_SNAKE_CASE_     ) == 0:
										__lowerCAmelCase				   =       default_version

					print(F"""Updating version to {version}."""     )
					global_version_update(SCREAMING_SNAKE_CASE_      ,      patch=SCREAMING_SNAKE_CASE_     )
					if not patch:
										print("Cleaning main README, don't forget to run `make fix-copies`."     )
										clean_main_ref_in_model_list()







def 	_a      (						):
					__lowerCAmelCase				   =       get_version()
					__lowerCAmelCase				   =       F"""{current_version.major}.{current_version.minor + 1}.0.dev0"""
					__lowerCAmelCase				   =       current_version.base_version

					# Check with the user we got that right.
					__lowerCAmelCase				   =       input(F"""Which version are we developing now? [{dev_version}]"""     )
					if len(SCREAMING_SNAKE_CASE_     ) == 0:
										__lowerCAmelCase				   =       dev_version

					print(F"""Updating version to {version}."""     )
					global_version_update(SCREAMING_SNAKE_CASE_     )
					print("Cleaning main README, don't forget to run `make fix-copies`."     )
					clean_main_ref_in_model_list()


if __name__ == "__main__":
					UpperCamelCase__    			=			argparse.ArgumentParser()
					parser.add_argument("""--post_release""", action="""store_true""", help="""Whether this is pre or post release.""")
					parser.add_argument("""--patch""", action="""store_true""", help="""Whether or not this is a patch release.""")
					UpperCamelCase__    			=			parser.parse_args()
					if not args.post_release:
										pre_release_work(patch=args.patch)
					elif args.patch:
										print("""Nothing to do after a patch :-)""")
					else:
										post_release_work()




