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

      __UpperCamelCase:  Dict						=						MODEL_FOR_CAUSAL_LM_MAPPING
      __UpperCamelCase:  Optional[Any]						=						TF_MODEL_FOR_CAUSAL_LM_MAPPING




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

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

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




      @require_tf
      def      _A	(       self						:				Union[str, Any]       ):
             _UpperCAmelCase     :  int							      =    pipeline(task="text-generation"						,					model="sshleifer/tiny-ctrl"						,					framework="tf"       )

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

             _UpperCAmelCase     :  List[str]							      =    text_generator(["This is a test", "This is a second test"]						,					do_sample=A       )
             self.assertEqual(
                 A						,					[
                     [
                         {
                             "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						:				Any						,					A						:				List[str]						,					A						:				List[str]						,					A						:				Union[str, Any]       ):
             _UpperCAmelCase     :  Optional[Any]							      =    TextGenerationPipeline(model=A						,					tokenizer=A       )
             return text_generator, ["This is a test", "Another test"]




      def      _A	(       self						:				Any       ):
             _UpperCAmelCase     :  int							      =    "Hello I believe in"
             _UpperCAmelCase     :  int							      =    pipeline("text-generation"						,					model="hf-internal-testing/tiny-random-gpt2"       )
             _UpperCAmelCase     :  Optional[int]							      =    text_generator(A       )
             self.assertEqual(
                 A						,					[{"generated_text": "Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe"}]						,					)

             _UpperCAmelCase     :  List[str]							      =    text_generator(A						,					stop_sequence=" fe"       )
             self.assertEqual(A						,					[{"generated_text": "Hello I believe in fe"}]       )




      def      _A	(       self						:				Optional[int]						,					A						:				Optional[Any]						,					A						:				List[Any]       ):
             _UpperCAmelCase     :  List[str]							      =    text_generator.model
             _UpperCAmelCase     :  str							      =    text_generator.tokenizer

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

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

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

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

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

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

             with self.assertRaises(A       ):
                    _UpperCAmelCase     :  List[str]							      =    text_generator("test"						,					return_full_text=A						,					return_text=A       )
             with self.assertRaises(A       ):
                    _UpperCAmelCase     :  Tuple							      =    text_generator("test"						,					return_full_text=A						,					return_tensors=A       )
             with self.assertRaises(A       ):
                    _UpperCAmelCase     :  List[str]							      =    text_generator("test"						,					return_text=A						,					return_tensors=A       )

             # 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__
             ):
                    _UpperCAmelCase     :  Any							      =    text_generator(""       )
                    self.assertEqual(A						,					[{"generated_text": ANY(A       )}]       )
             else:
                    with self.assertRaises((ValueError, AssertionError)       ):
                           _UpperCAmelCase     :  Dict							      =    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.
             _UpperCAmelCase     :  Any							      =    ["RwkvForCausalLM", "XGLMForCausalLM", "GPTNeoXForCausalLM"]
             if (
                 tokenizer.model_max_length < 10000
                 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       )

                    _UpperCAmelCase     :  Any							      =    text_generator("This is a test" * 500						,					handle_long_generation="hole"						,					max_new_tokens=20       )
                    # Hole strategy cannot work
                    with self.assertRaises(A       ):
                           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       ):
             import torch

             # Classic `model_kwargs`
             _UpperCAmelCase     :  Any							      =    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       )
             _UpperCAmelCase     :  List[str]							      =    pipe("This is a test"       )
             self.assertEqual(
                 A						,					[
                     {
                         "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.)
             _UpperCAmelCase     :  Tuple							      =    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       )
             _UpperCAmelCase     :  List[Any]							      =    pipe("This is a test"       )
             self.assertEqual(
                 A						,					[
                     {
                         "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
             _UpperCAmelCase     :  Tuple							      =    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       )
             _UpperCAmelCase     :  Optional[Any]							      =    pipe("This is a test"       )
             self.assertEqual(
                 A						,					[
                     {
                         "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						:				Tuple       ):
             import torch

             _UpperCAmelCase     :  Optional[int]							      =    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						:				Optional[Any]       ):
             import torch

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





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

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

             with CaptureLogger(A       ) as cl:
                    _UpperCAmelCase     :  List[str]							      =    text_generator(A						,					max_length=10       )
             self.assertNotIn(A						,					cl.out       )


'''simple docstring'''




import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
  from .tokenization_pegasus import PegasusTokenizer
else:
  __SCREAMING_SNAKE_CASE       :     Optional[Any]     =  None


__SCREAMING_SNAKE_CASE       :     Dict     =  logging.get_logger(__name__)


__SCREAMING_SNAKE_CASE       :     Union[str, Any]     =  """▁"""

__SCREAMING_SNAKE_CASE       :     str     =  {"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

__SCREAMING_SNAKE_CASE       :     int     =  {
    """vocab_file""": {"""google/pegasus-xsum""": """https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"""},
    """tokenizer_file""": {
        """google/pegasus-xsum""": """https://huggingface.co/google/pegasus-xsum/resolve/main/tokenizer.json"""
    },
}

__SCREAMING_SNAKE_CASE       :     str     =  {
    """google/pegasus-xsum""": 512,
}

class 	lowerCamelCase_	(snake_case__    ):
      '''simple docstring'''

      __UpperCamelCase:  Optional[int]						=						VOCAB_FILES_NAMES
      __UpperCamelCase:  Dict						=						PRETRAINED_VOCAB_FILES_MAP
      __UpperCamelCase:  List[Any]						=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __UpperCamelCase:  Optional[int]						=						PegasusTokenizer
      __UpperCamelCase:  Optional[Any]						=						["input_ids", "attention_mask"]




      def __init__(       self						:				Dict						,					A						:				List[str]=None						,					A						:				Union[str, Any]=None						,					A						:				Optional[int]="<pad>"						,					A						:				Tuple="</s>"						,					A						:				Union[str, Any]="<unk>"						,					A						:				Union[str, Any]="<mask_2>"						,					A						:				Dict="<mask_1>"						,					A						:				Union[str, Any]=None						,					A						:				int=103						,					**A						:				Optional[Any]						,					):
             _UpperCAmelCase     :  Dict							      =    offset

             if additional_special_tokens is not None:
                    if not isinstance(A						,					A       ):
                           raise TypeError(
                               F"""additional_special_tokens should be of type {type(A       )}, but is"""
                               F""" {type(A       )}"""       )

                    _UpperCAmelCase     :  Optional[int]							      =    (
                        ([mask_token_sent] + additional_special_tokens)
                        if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                        else additional_special_tokens
                    )
                    # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
                    additional_special_tokens_extended += [
                        F"""<unk_{i}>""" for i in range(len(A       )						,					self.offset - 1       )
                    ]

                    if len(set(A       )       ) != len(A       ):
                           raise ValueError(
                               "Please make sure that the provided additional_special_tokens do not contain an incorrectly"
                               F""" shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}."""       )
                    _UpperCAmelCase     :  Any							      =    additional_special_tokens_extended
             else:
                    _UpperCAmelCase     :  Dict							      =    [mask_token_sent] if mask_token_sent is not None else []
                    additional_special_tokens += [F"""<unk_{i}>""" for i in range(2						,					self.offset       )]

             super().__init__(
                 A						,					tokenizer_file=A						,					pad_token=A						,					eos_token=A						,					unk_token=A						,					mask_token=A						,					mask_token_sent=A						,					offset=A						,					additional_special_tokens=A						,					**A						,					)
             _UpperCAmelCase     :  Optional[Any]							      =    vocab_file
             _UpperCAmelCase     :  Optional[Any]							      =    False if not self.vocab_file else True




      def      _A	(       self						:				List[str]						,					A						:				Optional[Any]       ):
             _UpperCAmelCase     :  Any							      =    set(self.all_special_ids       )  # call it once instead of inside list comp
             all_special_ids.remove(self.unk_token_id       )  # <unk> is only sometimes special

             if all_special_ids != set(range(len(self.additional_special_tokens       ) + 3       )       ):
                    raise ValueError(
                        "There should be 3 special tokens: mask_token, pad_token, and eos_token +"
                        F""" {len(self.additional_special_tokens       )} additional_special_tokens, but got {all_special_ids}"""       )

             return [1 if x in all_special_ids else 0 for x in seq]




      def      _A	(       self						:				str						,					A						:				List						,					A						:				Optional[List] = None						,					A						:				bool = False       ):
             if already_has_special_tokens:
                    return self._special_token_mask(A       )
             elif token_ids_a is None:
                    return self._special_token_mask(A       ) + [1]
             else:
                    return self._special_token_mask(token_ids_a + token_ids_a       ) + [1]




      def      _A	(       self						:				Optional[int]						,					A						:				Union[str, Any]						,					A						:				int=None       ):
             if token_ids_a is None:
                    return token_ids_a + [self.eos_token_id]
             # We don't expect to process pairs, but leave the pair logic for API consistency
             return token_ids_a + token_ids_a + [self.eos_token_id]





      def      _A	(       self						:				Union[str, Any]						,					A						:				str						,					A						:				Optional[str] = None       ):
             if not self.can_save_slow_tokenizer:
                    raise ValueError(
                        "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
                        "tokenizer."       )

             if not os.path.isdir(A       ):
                    logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""       )
                    return
             _UpperCAmelCase     :  List[Any]							      =    os.path.join(
                 A						,					(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]       )

             if os.path.abspath(self.vocab_file       ) != os.path.abspath(A       ):
                    copyfile(self.vocab_file						,					A       )

             return (out_vocab_file,)

