
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"""







  lowerCamelCase__     =						MODEL_FOR_CAUSAL_LM_MAPPING
  lowerCamelCase__     =						TF_MODEL_FOR_CAUSAL_LM_MAPPING







  @require_torch
  def    __A     (							self :       Optional[int]   )  ->	Union[str, Any]:
        SCREAMING_SNAKE_CASE_          =  pipeline(task="text-generation"  ,     model="sshleifer/tiny-ctrl"  ,     framework="pt"   )
        # Using `do_sample=False` to force deterministic output
        SCREAMING_SNAKE_CASE_          =  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@@"
                    )
                }
            ]  ,     )

        SCREAMING_SNAKE_CASE_          =  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@@"
                        )
                    }
                ],
            ]  ,     )

        SCREAMING_SNAKE_CASE_          =  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__   )},
            ]  ,     )
        SCREAMING_SNAKE_CASE_          =  text_generator.model.config.eos_token_id
        SCREAMING_SNAKE_CASE_          =  "<pad>"
        SCREAMING_SNAKE_CASE_          =  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    __A     (							self :       List[Any]   )  ->	List[str]:
        SCREAMING_SNAKE_CASE_          =  pipeline(task="text-generation"  ,     model="sshleifer/tiny-ctrl"  ,     framework="tf"   )

        # Using `do_sample=False` to force deterministic output
        SCREAMING_SNAKE_CASE_          =  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,"
                    )
                }
            ]  ,     )

        SCREAMING_SNAKE_CASE_          =  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    __A     (							self :       Union[str, Any]  ,     __magic_name__ :       List[Any]  ,     __magic_name__ :       Dict  ,     __magic_name__ :       List[Any]   )  ->	Any:
        SCREAMING_SNAKE_CASE_          =  TextGenerationPipeline(model=lowerCamelCase__  ,     tokenizer=lowerCamelCase__   )
        return text_generator, ["This is a test", "Another test"]







  def    __A     (							self :       Tuple   )  ->	Union[str, Any]:
        SCREAMING_SNAKE_CASE_          =  "Hello I believe in"
        SCREAMING_SNAKE_CASE_          =  pipeline("text-generation"  ,     model="hf-internal-testing/tiny-random-gpt2"   )
        SCREAMING_SNAKE_CASE_          =  text_generator(lowerCamelCase__   )
        self.assertEqual(
            lowerCamelCase__  ,     [{"generated_text": "Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe"}]  ,     )

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







  def    __A     (							self :       Optional[Any]  ,     __magic_name__ :       Any  ,     __magic_name__ :       int   )  ->	List[Any]:
        SCREAMING_SNAKE_CASE_          =  text_generator.model
        SCREAMING_SNAKE_CASE_          =  text_generator.tokenizer

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

        SCREAMING_SNAKE_CASE_          =  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"]   )

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

        SCREAMING_SNAKE_CASE_          =  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"   )   )

        SCREAMING_SNAKE_CASE_          =  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:
              SCREAMING_SNAKE_CASE_          =  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__   ):
              SCREAMING_SNAKE_CASE_          =  text_generator("test"  ,     return_full_text=lowerCamelCase__  ,     return_text=lowerCamelCase__   )
        with self.assertRaises(lowerCamelCase__   ):
              SCREAMING_SNAKE_CASE_          =  text_generator("test"  ,     return_full_text=lowerCamelCase__  ,     return_tensors=lowerCamelCase__   )
        with self.assertRaises(lowerCamelCase__   ):
              SCREAMING_SNAKE_CASE_          =  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__
        ):
              SCREAMING_SNAKE_CASE_          =  text_generator(""   )
              self.assertEqual(lowerCamelCase__  ,     [{"generated_text": ANY(lowerCamelCase__   )}]   )
        else:
              with self.assertRaises((ValueError, AssertionError)   ):
                    SCREAMING_SNAKE_CASE_          =  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.
        SCREAMING_SNAKE_CASE_          =  ["RwkvForCausalLM", "XGLMForCausalLM", "GPTNeoXForCausalLM"]
        if (
            tokenizer.model_max_length < 10_000
            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   )

              SCREAMING_SNAKE_CASE_          =  text_generator("This is a test" * 500  ,     handle_long_generation="hole"  ,     max_new_tokens=20   )
              # Hole strategy cannot work
              with self.assertRaises(lowerCamelCase__   ):
                    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    __A     (							self :       str   )  ->	Optional[int]:
        import torch

        # Classic `model_kwargs`
        SCREAMING_SNAKE_CASE_          =  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   )
        SCREAMING_SNAKE_CASE_          =  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.)
        SCREAMING_SNAKE_CASE_          =  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   )
        SCREAMING_SNAKE_CASE_          =  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
        SCREAMING_SNAKE_CASE_          =  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   )
        SCREAMING_SNAKE_CASE_          =  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    __A     (							self :       Optional[Any]   )  ->	Any:
        import torch

        SCREAMING_SNAKE_CASE_          =  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 :       List[Any]   )  ->	Any:
        import torch

        SCREAMING_SNAKE_CASE_          =  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    __A     (							self :       Optional[int]   )  ->	List[Any]:
        SCREAMING_SNAKE_CASE_          =  "Hello world"
        SCREAMING_SNAKE_CASE_          =  pipeline("text-generation"  ,     model="hf-internal-testing/tiny-random-gpt2"   )
        if text_generator.model.framework == "tf":
              SCREAMING_SNAKE_CASE_          =  logging.get_logger("transformers.generation.tf_utils"   )
        else:
              SCREAMING_SNAKE_CASE_          =  logging.get_logger("transformers.generation.utils"   )
        SCREAMING_SNAKE_CASE_          =  "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:
              SCREAMING_SNAKE_CASE_          =  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:
              SCREAMING_SNAKE_CASE_          =  text_generator(lowerCamelCase__  ,     max_new_tokens=1   )
        self.assertNotIn(lowerCamelCase__  ,     cl.out   )

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



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


A     :  str					     =       {
    "configuration_rembert": ["REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "RemBertConfig", "RemBertOnnxConfig"]
}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A     :  Tuple					     =       ["RemBertTokenizer"]

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A     :  Tuple					     =       ["RemBertTokenizerFast"]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A     :  Tuple					     =       [
          "REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
          "RemBertForCausalLM",
          "RemBertForMaskedLM",
          "RemBertForMultipleChoice",
          "RemBertForQuestionAnswering",
          "RemBertForSequenceClassification",
          "RemBertForTokenClassification",
          "RemBertLayer",
          "RemBertModel",
          "RemBertPreTrainedModel",
          "load_tf_weights_in_rembert",
      ]


try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A     :  List[str]					     =       [
          "TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
          "TFRemBertForCausalLM",
          "TFRemBertForMaskedLM",
          "TFRemBertForMultipleChoice",
          "TFRemBertForQuestionAnswering",
          "TFRemBertForSequenceClassification",
          "TFRemBertForTokenClassification",
          "TFRemBertLayer",
          "TFRemBertModel",
          "TFRemBertPreTrainedModel",
      ]


if TYPE_CHECKING:
      from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_rembert import RemBertTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_rembert_fast import RemBertTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_rembert import (
                REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                RemBertForCausalLM,
                RemBertForMaskedLM,
                RemBertForMultipleChoice,
                RemBertForQuestionAnswering,
                RemBertForSequenceClassification,
                RemBertForTokenClassification,
                RemBertLayer,
                RemBertModel,
                RemBertPreTrainedModel,
                load_tf_weights_in_rembert,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_rembert import (
                TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFRemBertForCausalLM,
                TFRemBertForMaskedLM,
                TFRemBertForMultipleChoice,
                TFRemBertForQuestionAnswering,
                TFRemBertForSequenceClassification,
                TFRemBertForTokenClassification,
                TFRemBertLayer,
                TFRemBertModel,
                TFRemBertPreTrainedModel,
            )


else:
      import sys

      A     :  Tuple					     =       _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


