

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  lowercase				(       unittest.TestCase			):




				def     a							(					self      ,   snake_case      ,   snake_case       ):
											snake_case_							=    jnp.ones((batch_size, length)       ) / length
											return scores




				def     a							(					self       ):
											snake_case_							=    None
											snake_case_							=    20

											snake_case_							=    self._get_uniform_logits(batch_size=2      ,   length=snake_case       )

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

											# compute softmax
											snake_case_							=    jax.nn.softmax(snake_case      ,   axis=-1       )

											snake_case_							=    FlaxTemperatureLogitsWarper(temperature=0.5       )
											snake_case_							=    FlaxTemperatureLogitsWarper(temperature=1.3       )

											snake_case_							=    jax.nn.softmax(temp_dist_warper_sharper(snake_case      ,   scores.copy()      ,   cur_len=snake_case       )      ,   axis=-1       )
											snake_case_							=    jax.nn.softmax(temp_dist_warper_smoother(snake_case      ,   scores.copy()      ,   cur_len=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       ):
											snake_case_							=    None
											snake_case_							=    10
											snake_case_							=    2

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

											snake_case_							=    FlaxTopKLogitsWarper(3       )

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

											snake_case_							=    np.broadcast_to(np.arange(snake_case       )[None, :]      ,   (batch_size, length)       ).copy()
											snake_case_							=    top_k_warp_safety_check(snake_case      ,   snake_case      ,   cur_len=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       ):
											snake_case_							=    None
											snake_case_							=    10
											snake_case_							=    2

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

											snake_case_							=    FlaxTopPLogitsWarper(0.8       )
											snake_case_							=    np.exp(top_p_warp(snake_case      ,   snake_case      ,   cur_len=snake_case       )       )

											# dist should be filtered to keep min num values so that sum is >= top_p
											# exp (-inf) => 0
											snake_case_							=    np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]       )
											self.assertTrue(np.allclose(snake_case      ,   snake_case      ,   atol=1e-3       )       )

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

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

											# make sure at least 2 tokens are kept
											snake_case_							=    FlaxTopPLogitsWarper(0.9      ,   min_tokens_to_keep=2      ,   filter_value=0.0       )
											snake_case_							=    top_p_warp(snake_case      ,   snake_case      ,   cur_len=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       ):
											snake_case_							=    20
											snake_case_							=    4
											snake_case_							=    0

											snake_case_							=    FlaxMinLengthLogitsProcessor(min_length=10      ,   eos_token_id=snake_case       )

											# check that min length is applied at length 5
											snake_case_							=    ids_tensor((batch_size, 20)      ,   vocab_size=20       )
											snake_case_							=    5
											snake_case_							=    self._get_uniform_logits(snake_case      ,   snake_case       )
											snake_case_							=    min_dist_processor(snake_case      ,   snake_case      ,   cur_len=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
											snake_case_							=    self._get_uniform_logits(snake_case      ,   snake_case       )
											snake_case_							=    15
											snake_case_							=    min_dist_processor(snake_case      ,   snake_case      ,   cur_len=snake_case       )
											self.assertFalse(jnp.isinf(snake_case       ).any()       )




				def     a							(					self       ):
											snake_case_							=    20
											snake_case_							=    4
											snake_case_							=    0

											snake_case_							=    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case       )

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




				def     a							(					self       ):
											snake_case_							=    20
											snake_case_							=    4
											snake_case_							=    0
											snake_case_							=    5

											snake_case_							=    FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case      ,   eos_token_id=snake_case       )

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




				def     a							(					self       ):
											snake_case_							=    4
											snake_case_							=    10
											snake_case_							=    15
											snake_case_							=    2
											snake_case_							=    1
											snake_case_							=    15

											# dummy input_ids and scores
											snake_case_							=    ids_tensor((batch_size, sequence_length)      ,   snake_case       )
											snake_case_							=    input_ids.copy()

											snake_case_							=    self._get_uniform_logits(snake_case      ,   snake_case       )
											snake_case_							=    scores.copy()

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

											# instantiate all logits processors
											snake_case_							=    FlaxMinLengthLogitsProcessor(min_length=10      ,   eos_token_id=snake_case       )
											snake_case_							=    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case       )
											snake_case_							=    FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case      ,   eos_token_id=snake_case       )

											snake_case_							=    10

											# no processor list
											snake_case_							=    temp_dist_warp(snake_case      ,   snake_case      ,   cur_len=snake_case       )
											snake_case_							=    top_k_warp(snake_case      ,   snake_case      ,   cur_len=snake_case       )
											snake_case_							=    top_p_warp(snake_case      ,   snake_case      ,   cur_len=snake_case       )
											snake_case_							=    min_dist_proc(snake_case      ,   snake_case      ,   cur_len=snake_case       )
											snake_case_							=    bos_dist_proc(snake_case      ,   snake_case      ,   cur_len=snake_case       )
											snake_case_							=    eos_dist_proc(snake_case      ,   snake_case      ,   cur_len=snake_case       )

											# with processor list
											snake_case_							=    FlaxLogitsProcessorList(
											    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]       )
											snake_case_							=    processor(snake_case      ,   snake_case      ,   cur_len=snake_case       )

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

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




				def     a							(					self       ):
											snake_case_							=    4
											snake_case_							=    10
											snake_case_							=    15
											snake_case_							=    2
											snake_case_							=    1
											snake_case_							=    15

											# dummy input_ids and scores
											snake_case_							=    ids_tensor((batch_size, sequence_length)      ,   snake_case       )
											snake_case_							=    input_ids.copy()

											snake_case_							=    self._get_uniform_logits(snake_case      ,   snake_case       )
											snake_case_							=    scores.copy()

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

											# instantiate all logits processors
											snake_case_							=    FlaxMinLengthLogitsProcessor(min_length=10      ,   eos_token_id=snake_case       )
											snake_case_							=    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case       )
											snake_case_							=    FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case      ,   eos_token_id=snake_case       )

											snake_case_							=    10

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

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

											snake_case_							=    jax.jit(snake_case       )
											snake_case_							=    jax.jit(snake_case       )

											snake_case_							=    jitted_run_no_processor_list(snake_case      ,   snake_case      ,   snake_case       )
											snake_case_							=    jitted_run_processor_list(snake_case      ,   snake_case      ,   snake_case       )

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

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







import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser





def 					__lowerCamelCase	( UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    {}
							snake_case_							=    tokenizer(example['content']   ,							truncation=UpperCamelCase__      )['input_ids']
							snake_case_							=    len(example['content']      ) / len(output['input_ids']      )
							return output


_UpperCAmelCase		:   Dict										=       HfArgumentParser(PretokenizationArguments)
_UpperCAmelCase		:   List[Any]										=       parser.parse_args()
if args.num_workers is None:
					_UpperCAmelCase		:   Union[str, Any]										=       multiprocessing.cpu_count()
_UpperCAmelCase		:   int										=       AutoTokenizer.from_pretrained(args.tokenizer_dir)

_UpperCAmelCase		:   Optional[int]										=       time.time()
_UpperCAmelCase		:   List[str]										=       load_dataset(args.dataset_name, split="""train""")
print(F'''Dataset loaded in {time.time()-t_start:.2f}s''')

_UpperCAmelCase		:   Tuple										=       time.time()
_UpperCAmelCase		:   Union[str, Any]										=       ds.map(
    tokenize,
    num_proc=args.num_workers,
    remove_columns=[
        """repo_name""",
        """path""",
        """copies""",
        """size""",
        """content""",
        """license""",
        """hash""",
        """line_mean""",
        """line_max""",
        """alpha_frac""",
        """autogenerated""",
    ],
)
print(F'''Dataset tokenized in {time.time()-t_start:.2f}s''')

_UpperCAmelCase		:   Dict										=       time.time()
ds.push_to_hub(args.tokenized_data_repo)


print(F'''Data pushed to the hub in {time.time()-t_start:.2f}s''')




