





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     SCREAMING_SNAKE_CASE__ (						unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE							       =	MODEL_FOR_CAUSAL_LM_MAPPING
 __SCREAMING_SNAKE_CASE							       =	TF_MODEL_FOR_CAUSAL_LM_MAPPING



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

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

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



 @require_tf
 def 				UpperCamelCase			(							self ):
        A__											=						pipeline(task='''text-generation''',model='''sshleifer/tiny-ctrl''',framework='''tf''' )

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

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



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

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



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase ):
        A__											=						text_generator.model
        A__											=						text_generator.tokenizer

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

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

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

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

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

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

        with self.assertRaises(_UpperCAmelCase ):
               A__											=						text_generator('''test''',return_full_text=_UpperCAmelCase,return_text=_UpperCAmelCase )
        with self.assertRaises(_UpperCAmelCase ):
               A__											=						text_generator('''test''',return_full_text=_UpperCAmelCase,return_tensors=_UpperCAmelCase )
        with self.assertRaises(_UpperCAmelCase ):
               A__											=						text_generator('''test''',return_text=_UpperCAmelCase,return_tensors=_UpperCAmelCase )

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

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



 @require_torch
 @require_accelerate
 @require_torch_gpu
 def 				UpperCamelCase			(							self ):
        import torch

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

        A__											=						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 				UpperCamelCase			(							self ):
        import torch

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




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

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

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












'''simple docstring'''




from argparse import ArgumentParser

from datasets.commands.convert import ConvertCommand
from datasets.commands.dummy_data import DummyDataCommand
from datasets.commands.env import EnvironmentCommand
from datasets.commands.run_beam import RunBeamCommand
from datasets.commands.test import TestCommand
from datasets.utils.logging import set_verbosity_info






def        _lowerCAmelCase			(  __snake_case  : Tuple			)     ->     Dict:
     return {key.lstrip('-'			): value for key, value in zip(unknown_args[::2] ,							unknown_args[1::2]			)}






def        _lowerCAmelCase			(  )     ->     Tuple:
     __A						: int      =       ArgumentParser(
         'HuggingFace Datasets CLI tool' ,							usage='datasets-cli <command> [<args>]' ,							allow_abbrev=__snake_case			)
     __A						: Optional[Any]      =       parser.add_subparsers(help='datasets-cli command helpers'			)
     set_verbosity_info()

     # Register commands
     ConvertCommand.register_subcommand(__snake_case			)
     EnvironmentCommand.register_subcommand(__snake_case			)
     TestCommand.register_subcommand(__snake_case			)
     RunBeamCommand.register_subcommand(__snake_case			)
     DummyDataCommand.register_subcommand(__snake_case			)

     # Parse args
     __A ,__A						: Optional[Any]      =       parser.parse_known_args()
     if not hasattr(__snake_case ,							'func'			):
          parser.print_help()
          exit(1			)
     __A						: Any      =       parse_unknown_args(__snake_case			)

     # Run
     __A						: List[Any]      =       args.func(__snake_case ,							**__snake_case			)
     service.run()


if __name__ == "__main__":
  main()