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



	"""simple docstring"""




	def 		_lowercase    (			self    :     str			,				SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     int			)			->	Dict:
					lowercase_	    =   jnp.ones((batch_size, length)			) / length
					return scores
	def 		_lowercase    (			self    :     List[Any]			)			->	List[Any]:
					lowercase_	    =   None
					lowercase_	    =   2_0

					lowercase_	    =   self._get_uniform_logits(batch_size=2			,				length=SCREAMING_SNAKE_CASE_			)

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

					# compute softmax
					lowercase_	    =   jax.nn.softmax(SCREAMING_SNAKE_CASE_			,				axis=-1			)

					lowercase_	    =   FlaxTemperatureLogitsWarper(temperature=0.5			)
					lowercase_	    =   FlaxTemperatureLogitsWarper(temperature=1.3			)

					lowercase_	    =   jax.nn.softmax(temp_dist_warper_sharper(SCREAMING_SNAKE_CASE_			,				scores.copy()			,				cur_len=SCREAMING_SNAKE_CASE_			)			,				axis=-1			)
					lowercase_	    =   jax.nn.softmax(temp_dist_warper_smoother(SCREAMING_SNAKE_CASE_			,				scores.copy()			,				cur_len=SCREAMING_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 		_lowercase    (			self    :     Optional[Any]			)			->	List[str]:
					lowercase_	    =   None
					lowercase_	    =   1_0
					lowercase_	    =   2

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

					lowercase_	    =   FlaxTopKLogitsWarper(3			)

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

					lowercase_	    =   np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_			)[None, :]			,				(batch_size, length)			).copy()
					lowercase_	    =   top_k_warp_safety_check(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_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 		_lowercase    (			self    :     List[Any]			)			->	Union[str, Any]:
					lowercase_	    =   None
					lowercase_	    =   1_0
					lowercase_	    =   2

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

					lowercase_	    =   FlaxTopPLogitsWarper(0.8			)
					lowercase_	    =   np.exp(top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)			)

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

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

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

					# make sure at least 2 tokens are kept
					lowercase_	    =   FlaxTopPLogitsWarper(0.9			,				min_tokens_to_keep=2			,				filter_value=0.0			)
					lowercase_	    =   top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_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 		_lowercase    (			self    :     Tuple			)			->	Any:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0

					lowercase_	    =   FlaxMinLengthLogitsProcessor(min_length=1_0			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that min length is applied at length 5
					lowercase_	    =   ids_tensor((batch_size, 2_0)			,				vocab_size=2_0			)
					lowercase_	    =   5
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   min_dist_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_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
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   1_5
					lowercase_	    =   min_dist_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     List[str]			)			->	Dict:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0

					lowercase_	    =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that all scores are -inf except the bos_token_id score
					lowercase_	    =   ids_tensor((batch_size, 1)			,				vocab_size=2_0			)
					lowercase_	    =   1
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_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
					lowercase_	    =   3
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     int			)			->	Optional[int]:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0
					lowercase_	    =   5

					lowercase_	    =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that all scores are -inf except the eos_token_id when max_length is reached
					lowercase_	    =   ids_tensor((batch_size, 4)			,				vocab_size=2_0			)
					lowercase_	    =   4
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_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
					lowercase_	    =   3
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     int			)			->	Tuple:
					lowercase_	    =   4
					lowercase_	    =   1_0
					lowercase_	    =   1_5
					lowercase_	    =   2
					lowercase_	    =   1
					lowercase_	    =   1_5

					# dummy input_ids and scores
					lowercase_	    =   ids_tensor((batch_size, sequence_length)			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   input_ids.copy()

					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   scores.copy()

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

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

					lowercase_	    =   1_0

					# no processor list
					lowercase_	    =   temp_dist_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   top_k_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   min_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   bos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   eos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)

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

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

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




	def 		_lowercase    (			self    :     List[Any]			)			->	List[str]:
					lowercase_	    =   4
					lowercase_	    =   1_0
					lowercase_	    =   1_5
					lowercase_	    =   2
					lowercase_	    =   1
					lowercase_	    =   1_5

					# dummy input_ids and scores
					lowercase_	    =   ids_tensor((batch_size, sequence_length)			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   input_ids.copy()

					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   scores.copy()

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

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

					lowercase_	    =   1_0

					# no processor list
					def run_no_processor_list(SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     Optional[int]			,				SCREAMING_SNAKE_CASE_    :     int			):
									lowercase_	    =   temp_dist_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   top_k_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   min_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   bos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   eos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									return scores

					# with processor list
					def run_processor_list(SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     Dict			,				SCREAMING_SNAKE_CASE_    :     List[Any]			):
									lowercase_	    =   FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
									lowercase_	    =   processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									return scores

					lowercase_	    =   jax.jit(SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   jax.jit(SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   jitted_run_no_processor_list(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   jitted_run_processor_list(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)

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

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






import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY
@is_pipeline_test
@require_torch_or_tf
class 					lowercase__(							unittest.TestCase				):



	"""simple docstring"""




	a     :List[Any]       	=      MODEL_FOR_CAUSAL_LM_MAPPING
	a     :Tuple       	=      TF_MODEL_FOR_CAUSAL_LM_MAPPING
	@require_torch
	def 		_lowercase    (			self    :     Tuple			)			->	Dict:
					lowercase_	    =   pipeline(task='''text-generation'''			,				model='''sshleifer/tiny-ctrl'''			,				framework='''pt'''			)
					# Using `do_sample=False` to force deterministic output
					lowercase_	    =   text_generator('''This is a test'''			,				do_sample=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'''
					                ''' oscope. FiliFili@@'''
					            )
					        }
					    ]			,				)

					lowercase_	    =   text_generator(['''This is a test''', '''This is a second test''']			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        [
					            {
					                '''generated_text''': (
					                    '''This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'''
					                    ''' oscope. FiliFili@@'''
					                )
					            }
					        ],
					        [
					            {
					                '''generated_text''': (
					                    '''This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy'''
					                    ''' oscope. oscope. FiliFili@@'''
					                )
					            }
					        ],
					    ]			,				)

					lowercase_	    =   text_generator('''This is a test'''			,				do_sample=SCREAMING_SNAKE_CASE_			,				num_return_sequences=2			,				return_tensors=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					        {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					    ]			,				)
					lowercase_	    =   text_generator.model.config.eos_token_id
					lowercase_	    =   '''<pad>'''
					lowercase_	    =   text_generator(
					    ['''This is a test''', '''This is a second test''']			,				do_sample=SCREAMING_SNAKE_CASE_			,				num_return_sequences=2			,				batch_size=2			,				return_tensors=SCREAMING_SNAKE_CASE_			,				)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        [
					            {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					            {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					        ],
					        [
					            {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					            {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					        ],
					    ]			,				)
	@require_tf
	def 		_lowercase    (			self    :     str			)			->	Optional[int]:
					lowercase_	    =   pipeline(task='''text-generation'''			,				model='''sshleifer/tiny-ctrl'''			,				framework='''tf'''			)

					# Using `do_sample=False` to force deterministic output
					lowercase_	    =   text_generator('''This is a test'''			,				do_sample=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'''
					                ''' please,'''
					            )
					        }
					    ]			,				)

					lowercase_	    =   text_generator(['''This is a test''', '''This is a second test''']			,				do_sample=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        [
					            {
					                '''generated_text''': (
					                    '''This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'''
					                    ''' please,'''
					                )
					            }
					        ],
					        [
					            {
					                '''generated_text''': (
					                    '''This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes'''
					                    ''' Cannes 閲閲Cannes Cannes Cannes 攵 please,'''
					                )
					            }
					        ],
					    ]			,				)
	def 		_lowercase    (			self    :     Optional[Any]			,				SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     List[str]			)			->	Union[str, Any]:
					lowercase_	    =   TextGenerationPipeline(model=SCREAMING_SNAKE_CASE_			,				tokenizer=SCREAMING_SNAKE_CASE_			)
					return text_generator, ["This is a test", "Another test"]
	def 		_lowercase    (			self    :     List[Any]			)			->	Optional[int]:
					lowercase_	    =   '''Hello I believe in'''
					lowercase_	    =   pipeline('''text-generation'''			,				model='''hf-internal-testing/tiny-random-gpt2'''			)
					lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': '''Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'''}]			,				)

					lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			,				stop_sequence=''' fe'''			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': '''Hello I believe in fe'''}]			)
	def 		_lowercase    (			self    :     List[Any]			,				SCREAMING_SNAKE_CASE_    :     Tuple			,				SCREAMING_SNAKE_CASE_    :     List[str]			)			->	str:
					lowercase_	    =   text_generator.model
					lowercase_	    =   text_generator.tokenizer

					lowercase_	    =   text_generator('''This is a test'''			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					self.assertTrue(outputs[0]['''generated_text'''].startswith('''This is a test'''			)			)

					lowercase_	    =   text_generator('''This is a test'''			,				return_full_text=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					self.assertNotIn('''This is a test'''			,				outputs[0]['''generated_text''']			)

					lowercase_	    =   pipeline(task='''text-generation'''			,				model=SCREAMING_SNAKE_CASE_			,				tokenizer=SCREAMING_SNAKE_CASE_			,				return_full_text=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   text_generator('''This is a test'''			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					self.assertNotIn('''This is a test'''			,				outputs[0]['''generated_text''']			)

					lowercase_	    =   text_generator('''This is a test'''			,				return_full_text=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					self.assertTrue(outputs[0]['''generated_text'''].startswith('''This is a test'''			)			)

					lowercase_	    =   text_generator(['''This is great !''', '''Something else''']			,				num_return_sequences=2			,				do_sample=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        [{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}, {'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}],
					        [{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}, {'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}],
					    ]			,				)

					if text_generator.tokenizer.pad_token is not None:
									lowercase_	    =   text_generator(
									    ['''This is great !''', '''Something else''']			,				num_return_sequences=2			,				batch_size=2			,				do_sample=SCREAMING_SNAKE_CASE_			)
									self.assertEqual(
									    SCREAMING_SNAKE_CASE_			,				[
									        [{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}, {'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}],
									        [{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}, {'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}],
									    ]			,				)

					with self.assertRaises(SCREAMING_SNAKE_CASE_			):
									lowercase_	    =   text_generator('''test'''			,				return_full_text=SCREAMING_SNAKE_CASE_			,				return_text=SCREAMING_SNAKE_CASE_			)
					with self.assertRaises(SCREAMING_SNAKE_CASE_			):
									lowercase_	    =   text_generator('''test'''			,				return_full_text=SCREAMING_SNAKE_CASE_			,				return_tensors=SCREAMING_SNAKE_CASE_			)
					with self.assertRaises(SCREAMING_SNAKE_CASE_			):
									lowercase_	    =   text_generator('''test'''			,				return_text=SCREAMING_SNAKE_CASE_			,				return_tensors=SCREAMING_SNAKE_CASE_			)

					# Empty prompt is slighly special
					# it requires BOS token to exist.
					# Special case for Pegasus which will always append EOS so will
					# work even without BOS.
					if (
					    text_generator.tokenizer.bos_token_id is not None
					    or "Pegasus" in tokenizer.__class__.__name__
					    or "Git" in model.__class__.__name__
					):
									lowercase_	    =   text_generator(''''''			)
									self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					else:
									with self.assertRaises((ValueError, AssertionError)			):
													lowercase_	    =   text_generator(''''''			)

					if text_generator.framework == "tf":
									# TF generation does not support max_new_tokens, and it's impossible
									# to control long generation with only max_length without
									# fancy calculation, dismissing tests for now.
									return
					# We don't care about infinite range models.
					# They already work.
					# Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
					lowercase_	    =   ['''RwkvForCausalLM''', '''XGLMForCausalLM''', '''GPTNeoXForCausalLM''']
					if (
					    tokenizer.model_max_length < 1_0_0_0_0
					    and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
					):
									# Handling of large generations
									with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)			):
													text_generator('''This is a test''' * 5_0_0			,				max_new_tokens=2_0			)

									lowercase_	    =   text_generator('''This is a test''' * 5_0_0			,				handle_long_generation='''hole'''			,				max_new_tokens=2_0			)
									# Hole strategy cannot work
									with self.assertRaises(SCREAMING_SNAKE_CASE_			):
													text_generator(
													    '''This is a test''' * 5_0_0			,				handle_long_generation='''hole'''			,				max_new_tokens=tokenizer.model_max_length + 1_0			,				)
	@require_torch
	@require_accelerate
	@require_torch_gpu
	def 		_lowercase    (			self    :     Union[str, Any]			)			->	Tuple:
					import torch

					# Classic `model_kwargs`
					lowercase_	    =   pipeline(
					    model='''hf-internal-testing/tiny-random-bloom'''			,				model_kwargs={'''device_map''': '''auto''', '''torch_dtype''': torch.bfloataa}			,				)
					self.assertEqual(pipe.model.device			,				torch.device(0			)			)
					self.assertEqual(pipe.model.lm_head.weight.dtype			,				torch.bfloataa			)
					lowercase_	    =   pipe('''This is a test'''			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test test test test test test test test test test test test test test test test'''
					                ''' test'''
					            )
					        }
					    ]			,				)

					# Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
					lowercase_	    =   pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,				device_map='''auto'''			,				torch_dtype=torch.bfloataa			)
					self.assertEqual(pipe.model.device			,				torch.device(0			)			)
					self.assertEqual(pipe.model.lm_head.weight.dtype			,				torch.bfloataa			)
					lowercase_	    =   pipe('''This is a test'''			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test test test test test test test test test test test test test test test test'''
					                ''' test'''
					            )
					        }
					    ]			,				)

					# torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
					lowercase_	    =   pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,				device_map='''auto'''			)
					self.assertEqual(pipe.model.device			,				torch.device(0			)			)
					self.assertEqual(pipe.model.lm_head.weight.dtype			,				torch.floataa			)
					lowercase_	    =   pipe('''This is a test'''			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test test test test test test test test test test test test test test test test'''
					                ''' test'''
					            )
					        }
					    ]			,				)
	@require_torch
	@require_torch_gpu
	def 		_lowercase    (			self    :     Optional[int]			)			->	List[Any]:
					import torch

					lowercase_	    =   pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,				device=0			,				torch_dtype=torch.floataa			)
					pipe('''This is a test'''			)
	@require_torch
	@require_accelerate
	@require_torch_gpu
	def 		_lowercase    (			self    :     int			)			->	Dict:
					import torch

					lowercase_	    =   pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,				device_map='''auto'''			,				torch_dtype=torch.floataa			)
					pipe('''This is a test'''			,				do_sample=SCREAMING_SNAKE_CASE_			,				top_p=0.5			)




	def 		_lowercase    (			self    :     List[Any]			)			->	int:
					lowercase_	    =   '''Hello world'''
					lowercase_	    =   pipeline('''text-generation'''			,				model='''hf-internal-testing/tiny-random-gpt2'''			)
					if text_generator.model.framework == "tf":
									lowercase_	    =   logging.get_logger('''transformers.generation.tf_utils'''			)
					else:
									lowercase_	    =   logging.get_logger('''transformers.generation.utils'''			)
					lowercase_	    =   '''Both `max_new_tokens`'''  # The beggining of the message to be checked in this test

					# Both are set by the user -> log warning
					with CaptureLogger(SCREAMING_SNAKE_CASE_			) as cl:
									lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			,				max_length=1_0			,				max_new_tokens=1			)
					self.assertIn(SCREAMING_SNAKE_CASE_			,				cl.out			)

					# The user only sets one -> no warning
					with CaptureLogger(SCREAMING_SNAKE_CASE_			) as cl:
									lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			,				max_new_tokens=1			)
					self.assertNotIn(SCREAMING_SNAKE_CASE_			,				cl.out			)

					with CaptureLogger(SCREAMING_SNAKE_CASE_			) as cl:
									lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			,				max_length=1_0			)
					self.assertNotIn(SCREAMING_SNAKE_CASE_			,				cl.out			)




