
"""simple docstring"""





import os
from pathlib import Path

import numpy as np
import pytest
from pack_dataset import pack_data_dir
from parameterized import parameterized
from save_len_file import save_len_file
from torch.utils.data import DataLoader

from transformers import AutoTokenizer
from transformers.models.mbart.modeling_mbart import shift_tokens_right
from transformers.testing_utils import TestCasePlus, slow
from utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeqaSeqDataset, SeqaSeqDataset


_a								=						"""bert-base-cased"""
_a								=						"""google/pegasus-xsum"""
_a								=						[""" Sam ate lunch today.""", """Sams lunch ingredients."""]
_a								=						["""A very interesting story about what I ate for lunch.""", """Avocado, celery, turkey, coffee"""]
_a								=						"""patrickvonplaten/t5-tiny-random"""
_a								=						"""sshleifer/bart-tiny-random"""
_a								=						"""sshleifer/tiny-mbart"""
_a								=						"""sshleifer/tiny-marian-en-de"""



def  lowerCamelCase__						(     __snake_case,	__snake_case      )				->						Any:


					"""simple docstring"""





					_UpperCamelCase									=    '''\n'''.join(__snake_case      )
					Path(__snake_case      ).open('''w'''      ).writelines(__snake_case      )



def  lowerCamelCase__						(     __snake_case      )				->						Union[str, Any]:


					"""simple docstring"""





					for split in ["train", "val", "test"]:
										_dump_articles(os.path.join(__snake_case,	F'''{split}.source'''      ),	__snake_case      )
										_dump_articles(os.path.join(__snake_case,	F'''{split}.target'''      ),	__snake_case      )
					return tmp_dir






class        _UpperCAmelCase(       lowerCamelCase      ):



		@parameterized.expand(
		    [
		        MBART_TINY,
		        MARIAN_TINY,
		        T5_TINY,
		        BART_TINY,
		        PEGASUS_XSUM,
		    ]							,						)
		@slow
		def 		UpperCAmelCase				(   self							,						__a)			->			Dict:






							'''simple docstring'''
							_UpperCamelCase									=    AutoTokenizer.from_pretrained(__a)
							_UpperCamelCase									=    make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
							_UpperCamelCase									=    max(len(tokenizer.encode(__a)) for a in ARTICLES)
							_UpperCamelCase									=    max(len(tokenizer.encode(__a)) for a in SUMMARIES)
							_UpperCamelCase									=    4
							_UpperCamelCase									=    8
							assert max_len_target > max_src_len  # Will be truncated
							assert max_len_source > max_src_len  # Will be truncated
							_UpperCamelCase					,    _UpperCamelCase									=    '''ro_RO''', '''de_DE'''  # ignored for all but mbart, but never causes error.
							_UpperCamelCase									=    SeqaSeqDataset(
							    __a							,						data_dir=__a							,						type_path='''train'''							,						max_source_length=__a							,						max_target_length=__a							,						src_lang=__a							,						tgt_lang=__a							,						)
							_UpperCamelCase									=    DataLoader(__a							,						batch_size=2							,						collate_fn=train_dataset.collate_fn)
							for batch in dataloader:
												assert isinstance(__a							,						__a)
												assert batch["attention_mask"].shape == batch["input_ids"].shape
												# show that articles were trimmed.
												assert batch["input_ids"].shape[1] == max_src_len
												# show that targets are the same len
												assert batch["labels"].shape[1] == max_tgt_len
												if tok_name != MBART_TINY:
																	continue
												# check language codes in correct place
												_UpperCamelCase									=    shift_tokens_right(batch['''labels''']							,						tokenizer.pad_token_id)
												assert batch["decoder_input_ids"][0, 0].item() == tokenizer.lang_code_to_id[tgt_lang]
												assert batch["decoder_input_ids"][0, -1].item() == tokenizer.eos_token_id
												assert batch["input_ids"][0, -2].item() == tokenizer.eos_token_id
												assert batch["input_ids"][0, -1].item() == tokenizer.lang_code_to_id[src_lang]

												break  # No need to test every batch



		@parameterized.expand([BART_TINY, BERT_BASE_CASED])
		def 		UpperCAmelCase				(   self							,						__a)			->			Tuple:






							'''simple docstring'''
							_UpperCamelCase									=    AutoTokenizer.from_pretrained(__a)
							_UpperCamelCase									=    make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
							_UpperCamelCase									=    max(len(tokenizer.encode(__a)) for a in ARTICLES)
							_UpperCamelCase									=    max(len(tokenizer.encode(__a)) for a in SUMMARIES)
							_UpperCamelCase									=    4
							_UpperCamelCase									=    LegacySeqaSeqDataset(
							    __a							,						data_dir=__a							,						type_path='''train'''							,						max_source_length=20							,						max_target_length=__a							,						)
							_UpperCamelCase									=    DataLoader(__a							,						batch_size=2							,						collate_fn=train_dataset.collate_fn)
							for batch in dataloader:
												assert batch["attention_mask"].shape == batch["input_ids"].shape
												# show that articles were trimmed.
												assert batch["input_ids"].shape[1] == max_len_source
												assert 20 >= batch["input_ids"].shape[1]  # trimmed significantly
												# show that targets were truncated
												assert batch["labels"].shape[1] == trunc_target  # Truncated
												assert max_len_target > trunc_target  # Truncated
												break  # No need to test every batch



		def 		UpperCAmelCase				(   self)			->			Dict:






							'''simple docstring'''
							_UpperCamelCase									=    AutoTokenizer.from_pretrained('''facebook/mbart-large-cc25''')

							_UpperCamelCase									=    Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()))
							_UpperCamelCase									=    tmp_dir.joinpath('''train.source''').open().readlines()
							_UpperCamelCase									=    Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()))
							pack_data_dir(__a							,						__a							,						1_28							,						__a)
							_UpperCamelCase									=    {x.name for x in tmp_dir.iterdir()}
							_UpperCamelCase									=    {x.name for x in save_dir.iterdir()}
							_UpperCamelCase									=    save_dir.joinpath('''train.source''').open().readlines()
							# orig: [' Sam ate lunch today.\n', 'Sams lunch ingredients.']
							# desired_packed: [' Sam ate lunch today.\n Sams lunch ingredients.']
							assert len(__a) < len(__a)
							assert len(__a) == 1
							assert len(packed_examples[0]) == sum(len(__a) for x in orig_examples)
							assert orig_paths == new_paths



		@pytest.mark.skipif(not FAIRSEQ_AVAILABLE							,						reason='''This test requires fairseq''')
		def 		UpperCAmelCase				(   self)			->			str:






							'''simple docstring'''
							if not FAIRSEQ_AVAILABLE:
												return
							_UpperCamelCase					,    _UpperCamelCase					,    _UpperCamelCase									=    self._get_dataset(max_len=64)
							_UpperCamelCase									=    64
							_UpperCamelCase									=    ds.make_dynamic_sampler(__a							,						required_batch_size_multiple=__a)
							_UpperCamelCase									=    [len(__a) for x in batch_sampler]
							assert len(set(__a)) > 1  # it's not dynamic batch size if every batch is the same length
							assert sum(__a) == len(__a)  # no dropped or added examples
							_UpperCamelCase									=    DataLoader(__a							,						batch_sampler=__a							,						collate_fn=ds.collate_fn							,						num_workers=2)
							_UpperCamelCase									=    []
							_UpperCamelCase									=    []
							for batch in data_loader:
												_UpperCamelCase									=    batch['''input_ids'''].shape
												_UpperCamelCase									=    src_shape[0]
												assert bs % required_batch_size_multiple == 0 or bs < required_batch_size_multiple
												_UpperCamelCase									=    np.product(batch['''input_ids'''].shape)
												num_src_per_batch.append(__a)
												if num_src_tokens > (max_tokens * 1.1):
																	failures.append(__a)
							assert num_src_per_batch[0] == max(__a)
							if failures:
												raise AssertionError(F'''too many tokens in {len(__a)} batches''')



		def 		UpperCAmelCase				(   self)			->			Optional[Any]:






							'''simple docstring'''
							_UpperCamelCase					,    _UpperCamelCase					,    _UpperCamelCase									=    self._get_dataset(max_len=5_12)
							_UpperCamelCase									=    2
							_UpperCamelCase									=    ds.make_sortish_sampler(__a							,						shuffle=__a)

							_UpperCamelCase									=    DataLoader(__a							,						batch_size=__a							,						collate_fn=ds.collate_fn							,						num_workers=2)
							_UpperCamelCase									=    DataLoader(__a							,						batch_size=__a							,						collate_fn=ds.collate_fn							,						num_workers=2							,						sampler=__a)

							_UpperCamelCase									=    tokenizer.pad_token_id

							def count_pad_tokens(__a							,						__a="input_ids"):
												return [batch[k].eq(__a).sum().item() for batch in data_loader]

							assert sum(count_pad_tokens(__a							,						k='''labels''')) < sum(count_pad_tokens(__a							,						k='''labels'''))
							assert sum(count_pad_tokens(__a)) < sum(count_pad_tokens(__a))
							assert len(__a) == len(__a)



		def 		UpperCAmelCase				(   self							,						__a=10_00							,						__a=1_28)			->			int:






							'''simple docstring'''
							if os.getenv('''USE_REAL_DATA'''							,						__a):
												_UpperCamelCase									=    '''examples/seq2seq/wmt_en_ro'''
												_UpperCamelCase									=    max_len * 2 * 64
												if not Path(__a).joinpath('''train.len''').exists():
																	save_len_file(__a							,						__a)
							else:
												_UpperCamelCase									=    '''examples/seq2seq/test_data/wmt_en_ro'''
												_UpperCamelCase									=    max_len * 4
												save_len_file(__a							,						__a)

							_UpperCamelCase									=    AutoTokenizer.from_pretrained(__a)
							_UpperCamelCase									=    SeqaSeqDataset(
							    __a							,						data_dir=__a							,						type_path='''train'''							,						max_source_length=__a							,						max_target_length=__a							,						n_obs=__a							,						)
							return ds, max_tokens, tokenizer



		def 		UpperCAmelCase				(   self)			->			List[Any]:






							'''simple docstring'''
							_UpperCamelCase					,    _UpperCamelCase					,    _UpperCamelCase									=    self._get_dataset()
							_UpperCamelCase									=    set(DistributedSortishSampler(__a							,						2_56							,						num_replicas=2							,						rank=0							,						add_extra_examples=__a))
							_UpperCamelCase									=    set(DistributedSortishSampler(__a							,						2_56							,						num_replicas=2							,						rank=1							,						add_extra_examples=__a))
							assert idsa.intersection(__a) == set()




		@parameterized.expand(
		    [
		        MBART_TINY,
		        MARIAN_TINY,
		        T5_TINY,
		        BART_TINY,
		        PEGASUS_XSUM,
		    ]							,						)
		def 		UpperCAmelCase				(   self							,						__a)			->			Optional[int]:






							'''simple docstring'''
							_UpperCamelCase									=    AutoTokenizer.from_pretrained(__a							,						use_fast=__a)
							if tok_name == MBART_TINY:
												_UpperCamelCase									=    SeqaSeqDataset(
												    __a							,						data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())							,						type_path='''train'''							,						max_source_length=4							,						max_target_length=8							,						src_lang='''EN'''							,						tgt_lang='''FR'''							,						)
												_UpperCamelCase									=    train_dataset.dataset_kwargs
												assert "src_lang" in kwargs and "tgt_lang" in kwargs
							else:
												_UpperCamelCase									=    SeqaSeqDataset(
												    __a							,						data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())							,						type_path='''train'''							,						max_source_length=4							,						max_target_length=8							,						)
												_UpperCamelCase									=    train_dataset.dataset_kwargs
												assert "add_prefix_space" not in kwargs if tok_name != BART_TINY else "add_prefix_space" in kwargs
												assert len(__a) == 1 if tok_name == BART_TINY else len(__a) == 0












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






				def 		snake_case	(			self      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : int					  =      jnp.ones((batch_size, length) ) / length
							return scores






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      None
							SCREAMING_SNAKE_CASE_   : int					  =      20

							SCREAMING_SNAKE_CASE_   : int					  =      self._get_uniform_logits(batch_size=2      ,length=snake_case__ )

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

							# compute softmax
							SCREAMING_SNAKE_CASE_   : int					  =      jax.nn.softmax(snake_case__      ,axis=-1 )

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      FlaxTemperatureLogitsWarper(temperature=0.5 )
							SCREAMING_SNAKE_CASE_   : Dict					  =      FlaxTemperatureLogitsWarper(temperature=1.3 )

							SCREAMING_SNAKE_CASE_   : Any					  =      jax.nn.softmax(temp_dist_warper_sharper(snake_case__      ,scores.copy()      ,cur_len=snake_case__ )      ,axis=-1 )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      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 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Any					  =      None
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      10
							SCREAMING_SNAKE_CASE_   : int					  =      2

							# create ramp distribution
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      np.broadcast_to(np.arange(snake_case__ )[None, :]      ,(batch_size, vocab_size) ).copy()
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      ramp_logits[1:, : vocab_size // 2] + vocab_size

							SCREAMING_SNAKE_CASE_   : List[str]					  =      FlaxTopKLogitsWarper(3 )

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

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      np.broadcast_to(np.arange(snake_case__ )[None, :]      ,(batch_size, length) ).copy()
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      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 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      None
							SCREAMING_SNAKE_CASE_   : str					  =      10
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      2

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

							SCREAMING_SNAKE_CASE_   : Dict					  =      FlaxTopPLogitsWarper(0.8 )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      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
							SCREAMING_SNAKE_CASE_   : str					  =      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
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      np.broadcast_to(np.arange(snake_case__ )[None, :]      ,(batch_size, vocab_size) ).copy() - (
							    vocab_size // 2
							)

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

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

							SCREAMING_SNAKE_CASE_   : str					  =      FlaxMinLengthLogitsProcessor(min_length=10      ,eos_token_id=snake_case__ )

							# check that min length is applied at length 5
							SCREAMING_SNAKE_CASE_   : Dict					  =      ids_tensor((batch_size, 20)      ,vocab_size=20 )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      5
							SCREAMING_SNAKE_CASE_   : List[str]					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Any					  =      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
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      15
							SCREAMING_SNAKE_CASE_   : Any					  =      min_dist_processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							self.assertFalse(jnp.isinf(snake_case__ ).any() )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      20
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      4
							SCREAMING_SNAKE_CASE_   : Dict					  =      0

							SCREAMING_SNAKE_CASE_   : int					  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )

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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      20
							SCREAMING_SNAKE_CASE_   : Any					  =      4
							SCREAMING_SNAKE_CASE_   : int					  =      0
							SCREAMING_SNAKE_CASE_   : Tuple					  =      5

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      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
							SCREAMING_SNAKE_CASE_   : Any					  =      ids_tensor((batch_size, 4)      ,vocab_size=20 )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      4
							SCREAMING_SNAKE_CASE_   : Dict					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      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
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      3
							SCREAMING_SNAKE_CASE_   : int					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      logits_processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							self.assertFalse(jnp.isinf(snake_case__ ).any() )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      4
							SCREAMING_SNAKE_CASE_   : Dict					  =      10
							SCREAMING_SNAKE_CASE_   : int					  =      15
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      2
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      1
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      15

							# dummy input_ids and scores
							SCREAMING_SNAKE_CASE_   : Any					  =      ids_tensor((batch_size, sequence_length)      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      input_ids.copy()

							SCREAMING_SNAKE_CASE_   : str					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      scores.copy()

							# instantiate all dist processors
							SCREAMING_SNAKE_CASE_   : int					  =      FlaxTemperatureLogitsWarper(temperature=0.5 )
							SCREAMING_SNAKE_CASE_   : int					  =      FlaxTopKLogitsWarper(3 )
							SCREAMING_SNAKE_CASE_   : Dict					  =      FlaxTopPLogitsWarper(0.8 )

							# instantiate all logits processors
							SCREAMING_SNAKE_CASE_   : Any					  =      FlaxMinLengthLogitsProcessor(min_length=10      ,eos_token_id=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Any					  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,eos_token_id=snake_case__ )

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      10

							# no processor list
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      temp_dist_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[str]					  =      top_k_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      top_p_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : str					  =      min_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      bos_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Tuple					  =      eos_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )

							# with processor list
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      FlaxLogitsProcessorList(
							    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
							SCREAMING_SNAKE_CASE_   : Any					  =      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 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      4
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      10
							SCREAMING_SNAKE_CASE_   : Dict					  =      15
							SCREAMING_SNAKE_CASE_   : Dict					  =      2
							SCREAMING_SNAKE_CASE_   : Tuple					  =      1
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      15

							# dummy input_ids and scores
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      ids_tensor((batch_size, sequence_length)      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      input_ids.copy()

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      scores.copy()

							# instantiate all dist processors
							SCREAMING_SNAKE_CASE_   : str					  =      FlaxTemperatureLogitsWarper(temperature=0.5 )
							SCREAMING_SNAKE_CASE_   : List[str]					  =      FlaxTopKLogitsWarper(3 )
							SCREAMING_SNAKE_CASE_   : int					  =      FlaxTopPLogitsWarper(0.8 )

							# instantiate all logits processors
							SCREAMING_SNAKE_CASE_   : Tuple					  =      FlaxMinLengthLogitsProcessor(min_length=10      ,eos_token_id=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )
							SCREAMING_SNAKE_CASE_   : str					  =      FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,eos_token_id=snake_case__ )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      10

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

							SCREAMING_SNAKE_CASE_   : Tuple					  =      jax.jit(snake_case__ )
							SCREAMING_SNAKE_CASE_   : Any					  =      jax.jit(snake_case__ )

							SCREAMING_SNAKE_CASE_   : int					  =      jitted_run_no_processor_list(snake_case__      ,snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[str]					  =      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() )



