


'''simple docstring'''






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





							'''simple docstring'''






							lowercase__     :  List[Any]      		= MODEL_FOR_CAUSAL_LM_MAPPING
							lowercase__     :  str      		= TF_MODEL_FOR_CAUSAL_LM_MAPPING







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

														lowerCAmelCase__			  =      text_generator(['''This is a test''', '''This is a second test''']  )
														self.assertEqual(
														    lowerCamelCase_   ,       [
														        [
														            {
														                '''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@@'''
														                )
														            }
														        ],
														    ]   ,       )

														lowerCAmelCase__			  =      text_generator('''This is a test'''   ,       do_sample=lowerCamelCase_   ,       num_return_sequences=2   ,       return_tensors=lowerCamelCase_  )
														self.assertEqual(
														    lowerCamelCase_   ,       [
														        {'''generated_token_ids''': ANY(lowerCamelCase_  )},
														        {'''generated_token_ids''': ANY(lowerCamelCase_  )},
														    ]   ,       )
														lowerCAmelCase__			  =      text_generator.model.config.eos_token_id
														lowerCAmelCase__			  =      '''<pad>'''
														lowerCAmelCase__			  =      text_generator(
														    ['''This is a test''', '''This is a second test''']   ,       do_sample=lowerCamelCase_   ,       num_return_sequences=2   ,       batch_size=2   ,       return_tensors=lowerCamelCase_   ,       )
														self.assertEqual(
														    lowerCamelCase_   ,       [
														        [
														            {'''generated_token_ids''': ANY(lowerCamelCase_  )},
														            {'''generated_token_ids''': ANY(lowerCamelCase_  )},
														        ],
														        [
														            {'''generated_token_ids''': ANY(lowerCamelCase_  )},
														            {'''generated_token_ids''': ANY(lowerCamelCase_  )},
														        ],
														    ]   ,       )







							@require_tf
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														lowerCAmelCase__			  =      pipeline(task='''text-generation'''   ,       model='''sshleifer/tiny-ctrl'''   ,       framework='''tf'''  )

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

														lowerCAmelCase__			  =      text_generator(['''This is a test''', '''This is a second test''']   ,       do_sample=lowerCamelCase_  )
														self.assertEqual(
														    lowerCamelCase_   ,       [
														        [
														            {
														                '''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     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Tuple:
														lowerCAmelCase__			  =      TextGenerationPipeline(model=lowerCamelCase_   ,       tokenizer=lowerCamelCase_  )
														return text_generator, ["This is a test", "Another test"]







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      '''Hello I believe in'''
														lowerCAmelCase__			  =      pipeline('''text-generation'''   ,       model='''hf-internal-testing/tiny-random-gpt2'''  )
														lowerCAmelCase__			  =      text_generator(lowerCamelCase_  )
														self.assertEqual(
														    lowerCamelCase_   ,       [{'''generated_text''': '''Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'''}]   ,       )

														lowerCAmelCase__			  =      text_generator(lowerCamelCase_   ,       stop_sequence=''' fe'''  )
														self.assertEqual(lowerCamelCase_   ,       [{'''generated_text''': '''Hello I believe in fe'''}]  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Optional[Any]:
														lowerCAmelCase__			  =      text_generator.model
														lowerCAmelCase__			  =      text_generator.tokenizer

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

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

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

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

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

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

														with self.assertRaises(lowerCamelCase_  ):
																					lowerCAmelCase__			  =      text_generator('''test'''   ,       return_full_text=lowerCamelCase_   ,       return_text=lowerCamelCase_  )
														with self.assertRaises(lowerCamelCase_  ):
																					lowerCAmelCase__			  =      text_generator('''test'''   ,       return_full_text=lowerCamelCase_   ,       return_tensors=lowerCamelCase_  )
														with self.assertRaises(lowerCamelCase_  ):
																					lowerCAmelCase__			  =      text_generator('''test'''   ,       return_text=lowerCamelCase_   ,       return_tensors=lowerCamelCase_  )

														# 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__
														):
																					lowerCAmelCase__			  =      text_generator(''''''  )
																					self.assertEqual(lowerCamelCase_   ,       [{'''generated_text''': ANY(lowerCamelCase_  )}]  )
														else:
																					with self.assertRaises((ValueError, AssertionError)  ):
																												lowerCAmelCase__			  =      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.
														lowerCAmelCase__			  =      ['''RwkvForCausalLM''', '''XGLMForCausalLM''', '''GPTNeoXForCausalLM''']
														if (
														    tokenizer.model_max_length < 1_00_00
														    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_00   ,       max_new_tokens=20  )

																					lowerCAmelCase__			  =      text_generator('''This is a test''' * 5_00   ,       handle_long_generation='''hole'''   ,       max_new_tokens=20  )
																					# Hole strategy cannot work
																					with self.assertRaises(lowerCamelCase_  ):
																												text_generator(
																												    '''This is a test''' * 5_00   ,       handle_long_generation='''hole'''   ,       max_new_tokens=tokenizer.model_max_length + 10   ,       )







							@require_torch
							@require_accelerate
							@require_torch_gpu
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														import torch

														# Classic `model_kwargs`
														lowerCAmelCase__			  =      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  )
														lowerCAmelCase__			  =      pipe('''This is a test'''  )
														self.assertEqual(
														    lowerCamelCase_   ,       [
														        {
														            '''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.)
														lowerCAmelCase__			  =      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  )
														lowerCAmelCase__			  =      pipe('''This is a test'''  )
														self.assertEqual(
														    lowerCamelCase_   ,       [
														        {
														            '''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
														lowerCAmelCase__			  =      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  )
														lowerCAmelCase__			  =      pipe('''This is a test'''  )
														self.assertEqual(
														    lowerCamelCase_   ,       [
														        {
														            '''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     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														import torch

														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														import torch

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														lowerCAmelCase__			  =      '''Hello world'''
														lowerCAmelCase__			  =      pipeline('''text-generation'''   ,       model='''hf-internal-testing/tiny-random-gpt2'''  )
														if text_generator.model.framework == "tf":
																					lowerCAmelCase__			  =      logging.get_logger('''transformers.generation.tf_utils'''  )
														else:
																					lowerCAmelCase__			  =      logging.get_logger('''transformers.generation.utils'''  )
														lowerCAmelCase__			  =      '''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(lowerCamelCase_  ) as cl:
																					lowerCAmelCase__			  =      text_generator(lowerCamelCase_   ,       max_length=10   ,       max_new_tokens=1  )
														self.assertIn(lowerCamelCase_   ,       cl.out  )

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

														with CaptureLogger(lowerCamelCase_  ) as cl:
																					lowerCAmelCase__			  =      text_generator(lowerCamelCase_   ,       max_length=10  )
														self.assertNotIn(lowerCamelCase_   ,       cl.out  )






'''simple docstring'''

import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
  import torch

  from transformers.generation import (
      MaxLengthCriteria,
      MaxNewTokensCriteria,
      MaxTimeCriteria,
      StoppingCriteriaList,
      validate_stopping_criteria,
  )






@require_torch
class   UpperCAmelCase      (			unittest.TestCase							):
 def 	UpperCAmelCase_				(							self :int		,			lowercase_ :str				)->				Any:
     A__							=     3
     A__							=     2_50

     A__							=     ids_tensor((batch_size, length)		,			lowercase_				)
     A__							=     torch.ones((batch_size, length)		,			device=lowercase_		,			dtype=torch.float				) / length
     return input_ids, scores
 def 	UpperCAmelCase_				(							self :str				)->				Dict:
     A__,      A__							=     self._get_tensors(5				)

     A__							=     StoppingCriteriaList(
         [
             MaxLengthCriteria(max_length=10				),
             MaxTimeCriteria(max_time=0.1				),
         ]				)

     self.assertFalse(criteria(lowercase_		,			lowercase_				)				)

     A__,      A__							=     self._get_tensors(9				)
     self.assertFalse(criteria(lowercase_		,			lowercase_				)				)

     A__,      A__							=     self._get_tensors(10				)
     self.assertTrue(criteria(lowercase_		,			lowercase_				)				)
 def 	UpperCAmelCase_				(							self :List[Any]				)->				List[Any]:
     A__							=     MaxLengthCriteria(max_length=10				)

     A__,      A__							=     self._get_tensors(5				)
     self.assertFalse(criteria(lowercase_		,			lowercase_				)				)

     A__,      A__							=     self._get_tensors(9				)
     self.assertFalse(criteria(lowercase_		,			lowercase_				)				)

     A__,      A__							=     self._get_tensors(10				)
     self.assertTrue(criteria(lowercase_		,			lowercase_				)				)
 def 	UpperCAmelCase_				(							self :Tuple				)->				int:
     A__							=     MaxNewTokensCriteria(start_length=5		,			max_new_tokens=5				)

     A__,      A__							=     self._get_tensors(5				)
     self.assertFalse(criteria(lowercase_		,			lowercase_				)				)

     A__,      A__							=     self._get_tensors(9				)
     self.assertFalse(criteria(lowercase_		,			lowercase_				)				)

     A__,      A__							=     self._get_tensors(10				)
     self.assertTrue(criteria(lowercase_		,			lowercase_				)				)

     A__							=     StoppingCriteriaList([criteria]				)
     self.assertEqual(criteria_list.max_length		,			10				)
 def 	UpperCAmelCase_				(							self :Dict				)->				int:
     A__,      A__							=     self._get_tensors(5				)

     A__							=     MaxTimeCriteria(max_time=0.1				)
     self.assertFalse(criteria(lowercase_		,			lowercase_				)				)

     A__							=     MaxTimeCriteria(max_time=0.1		,			initial_timestamp=time.time() - 0.2				)
     self.assertTrue(criteria(lowercase_		,			lowercase_				)				)



 def 	UpperCAmelCase_				(							self :Union[str, Any]				)->				str:
     validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10				)]				)		,			10				)

     with self.assertWarns(lowercase_				):
         validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10				)]				)		,			11				)

     A__							=     validate_stopping_criteria(StoppingCriteriaList()		,			11				)

     self.assertEqual(len(lowercase_				)		,			1				)



