


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





  """simple docstring"""



  a_							:str				       =MODEL_FOR_CAUSAL_LM_MAPPING
  a_							:Optional[Any]				       =TF_MODEL_FOR_CAUSAL_LM_MAPPING

  @require_torch
  def __a    (       self  :				Tuple    ):





       '''simple docstring'''

       __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=snake_case_    )
       self.assertEqual(
           snake_case_					,				[
               {
                   """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(
           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@@"""
                       )
                   }
               ],
           ]					,				)

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

  @require_tf
  def __a    (       self  :				Optional[Any]    ):





       '''simple docstring'''

       __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=snake_case_    )
       self.assertEqual(
           snake_case_					,				[
               {
                   """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=snake_case_    )
       self.assertEqual(
           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 __a    (       self  :				str					,				SCREAMING_SNAKE_CASE__  :				Dict					,				SCREAMING_SNAKE_CASE__  :				Optional[Any]					,				SCREAMING_SNAKE_CASE__  :				Union[str, Any]    ):





       '''simple docstring'''

       __a   			= TextGenerationPipeline(model=snake_case_					,				tokenizer=snake_case_    )
       return text_generator, ["This is a test", "Another test"]

  def __a    (       self  :				Tuple    ):





       '''simple docstring'''

       __a   			= """Hello I believe in"""
       __a   			= pipeline("""text-generation"""					,				model="""hf-internal-testing/tiny-random-gpt2"""    )
       __a   			= text_generator(snake_case_    )
       self.assertEqual(
           snake_case_					,				[{"""generated_text""": """Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe"""}]					,				)

       __a   			= text_generator(snake_case_					,				stop_sequence=""" fe"""    )
       self.assertEqual(snake_case_					,				[{"""generated_text""": """Hello I believe in fe"""}]    )

  def __a    (       self  :				Union[str, Any]					,				SCREAMING_SNAKE_CASE__  :				Union[str, Any]					,				SCREAMING_SNAKE_CASE__  :				Dict    ):





       '''simple docstring'''

       __a   			= text_generator.model
       __a   			= text_generator.tokenizer

       __a   			= text_generator("""This is a test"""    )
       self.assertEqual(snake_case_					,				[{"""generated_text""": ANY(snake_case_    )}]    )
       self.assertTrue(outputs[0]["""generated_text"""].startswith("""This is a test"""    )    )

       __a   			= text_generator("""This is a test"""					,				return_full_text=snake_case_    )
       self.assertEqual(snake_case_					,				[{"""generated_text""": ANY(snake_case_    )}]    )
       self.assertNotIn("""This is a test"""					,				outputs[0]["""generated_text"""]    )

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

       __a   			= text_generator("""This is a test"""					,				return_full_text=snake_case_    )
       self.assertEqual(snake_case_					,				[{"""generated_text""": ANY(snake_case_    )}]    )
       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=snake_case_    )
       self.assertEqual(
           snake_case_					,				[
               [{"""generated_text""": ANY(snake_case_    )}, {"""generated_text""": ANY(snake_case_    )}],
               [{"""generated_text""": ANY(snake_case_    )}, {"""generated_text""": ANY(snake_case_    )}],
           ]					,				)

       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=snake_case_    )
            self.assertEqual(
                snake_case_					,				[
                    [{"""generated_text""": ANY(snake_case_    )}, {"""generated_text""": ANY(snake_case_    )}],
                    [{"""generated_text""": ANY(snake_case_    )}, {"""generated_text""": ANY(snake_case_    )}],
                ]					,				)

       with self.assertRaises(snake_case_    ):
            __a   			= text_generator("""test"""					,				return_full_text=snake_case_					,				return_text=snake_case_    )
       with self.assertRaises(snake_case_    ):
            __a   			= text_generator("""test"""					,				return_full_text=snake_case_					,				return_tensors=snake_case_    )
       with self.assertRaises(snake_case_    ):
            __a   			= text_generator("""test"""					,				return_text=snake_case_					,				return_tensors=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__
       ):
            __a   			= text_generator(""""""    )
            self.assertEqual(snake_case_					,				[{"""generated_text""": ANY(snake_case_    )}]    )
       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_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    )

            __a   			= 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(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 __a    (       self  :				Any    ):





       '''simple docstring'''

       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(
           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.)
       __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(
           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
       __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(
           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 __a    (       self  :				List[str]    ):





       '''simple docstring'''

       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 __a    (       self  :				Any    ):





       '''simple docstring'''

       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=snake_case_					,				top_p=0.5    )


  def __a    (       self  :				Optional[int]    ):





       '''simple docstring'''

       __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(snake_case_    ) as cl:
            __a   			= text_generator(snake_case_					,				max_length=1_0					,				max_new_tokens=1    )
       self.assertIn(snake_case_					,				cl.out    )

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

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








from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowerCamelCase_    :			Optional[int]    =	{
    """configuration_graphormer""": ["""GRAPHORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""", """GraphormerConfig"""],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCamelCase_    :			Union[str, Any]    =	[
							    """GRAPHORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """GraphormerForGraphClassification""",
							    """GraphormerModel""",
							    """GraphormerPreTrainedModel""",
							]


if TYPE_CHECKING:
							from .configuration_graphormer import GRAPHORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, GraphormerConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_graphormer import (
														    GRAPHORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
														    GraphormerForGraphClassification,
														    GraphormerModel,
														    GraphormerPreTrainedModel,
														)


else:
							import sys

							lowerCamelCase_    :			Optional[Any]    =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






