


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




 __lowerCAmelCase			 =  MODEL_FOR_CAUSAL_LM_MAPPING
 __lowerCAmelCase			 =  TF_MODEL_FOR_CAUSAL_LM_MAPPING


 @require_torch
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[int]:
     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=__A						)
     self.assertEqual(
         __A			,   [
             {
                 '''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(
         __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@@'''
                     )
                 }
             ],
         ]			,   )

     a					=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						)},
         ]			,   )
     a					=text_generator.model.config.eos_token_id
     a					='''<pad>'''
     a					=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 				SCREAMING_SNAKE_CASE							(		self						) ->					List[str]:
     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=__A						)
     self.assertEqual(
         __A			,   [
             {
                 '''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=__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 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A						) ->					List[str]:
     a					=TextGenerationPipeline(model=__A			,   tokenizer=__A						)
     return text_generator, ["This is a test", "Another test"]


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

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


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A						) ->					str:
     a					=text_generator.model
     a					=text_generator.tokenizer

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

     a					=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''']						)

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

     a					=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'''						)						)

     a					=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:
         a					=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						):
         a					=text_generator('''test'''			,   return_full_text=__A			,   return_text=__A						)
     with self.assertRaises(__A						):
         a					=text_generator('''test'''			,   return_full_text=__A			,   return_tensors=__A						)
     with self.assertRaises(__A						):
         a					=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__
     ):
         a					=text_generator(''''''						)
         self.assertEqual(__A			,   [{'''generated_text''': ANY(__A						)}]						)
     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(__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 				SCREAMING_SNAKE_CASE							(		self						) ->					List[str]:
     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(
         __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.)
     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(
         __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
     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(
         __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 				SCREAMING_SNAKE_CASE							(		self						) ->					List[str]:
     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 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     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=__A			,   top_p=0.5						)





 def 				SCREAMING_SNAKE_CASE							(		self						) ->					int:
     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(__A						) as cl:
         a					=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:
         a					=text_generator(__A			,   max_new_tokens=1						)
     self.assertNotIn(__A			,   cl.out						)

     with CaptureLogger(__A						) as cl:
         a					=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 import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
       from .tokenization_rembert import RemBertTokenizer
else:
       lowerCamelCase_       :       Optional[int]      						=  None

lowerCamelCase_       :       Union[str, Any]      						=  logging.get_logger(__name__)
lowerCamelCase_       :       str      						=  {"""vocab_file""": """sentencepiece.model""", """tokenizer_file""": """tokenizer.json"""}

lowerCamelCase_       :       int      						=  {
    """vocab_file""": {
        """google/rembert""": """https://huggingface.co/google/rembert/resolve/main/sentencepiece.model""",
    },
    """tokenizer_file""": {
        """google/rembert""": """https://huggingface.co/google/rembert/resolve/main/tokenizer.json""",
    },
}

lowerCamelCase_       :       Optional[Any]      						=  {
    """google/rembert""": 2_5_6,
}

lowerCamelCase_       :       Optional[Any]      						=  """▁"""






class 				__A    (	_SCREAMING_SNAKE_CASE      ):






 """simple docstring"""




 __lowerCAmelCase			 =  VOCAB_FILES_NAMES
 __lowerCAmelCase			 =  PRETRAINED_VOCAB_FILES_MAP
 __lowerCAmelCase			 =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCAmelCase			 =  RemBertTokenizer


 def __init__(		self			,   __A=None			,   __A=None			,   __A=True			,   __A=True			,   __A=False			,   __A="[CLS]"			,   __A="[SEP]"			,   __A="<unk>"			,   __A="[SEP]"			,   __A="<pad>"			,   __A="[CLS]"			,   __A="[MASK]"			,   **__A			,   ) ->					Dict:
     # Mask token behave like a normal word, i.e. include the space before it
     a					=AddedToken(__A			,   lstrip=__A			,   rstrip=__A						) if isinstance(__A			,   __A						) else mask_token

     super().__init__(
         __A			,   tokenizer_file=__A			,   do_lower_case=__A			,   remove_space=__A			,   keep_accents=__A			,   bos_token=__A			,   eos_token=__A			,   unk_token=__A			,   sep_token=__A			,   pad_token=__A			,   cls_token=__A			,   mask_token=__A			,   **__A			,   )

     a					=do_lower_case
     a					=remove_space
     a					=keep_accents
     a					=vocab_file
     a					=False if not self.vocab_file else True


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A = None						) ->					List[int]:
     a					=[self.sep_token_id]
     a					=[self.cls_token_id]
     if token_ids_a is None:
         return cls + token_ids_a + sep
     return cls + token_ids_a + sep + token_ids_a + sep


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A = None			,   __A = False						) ->					List[int]:

     if already_has_special_tokens:
         if token_ids_a is not None:
             raise ValueError(
                 '''You should not supply a second sequence if the provided sequence of '''
                 '''ids is already formatted with special tokens for the model.'''						)
         return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

     if token_ids_a is not None:
         return [1] + ([0] * len(__A						)) + [1] + ([0] * len(__A						)) + [1]
     return [1] + ([0] * len(__A						)) + [1]


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A = None						) ->					List[int]:
     a					=[self.sep_token_id]
     a					=[self.cls_token_id]

     if token_ids_a is None:
         return len(cls + token_ids_a + sep						) * [0]
     return len(cls + token_ids_a + sep						) * [0] + len(token_ids_a + sep						) * [1]





 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A = None						) ->					Tuple[str]:
     if not os.path.isdir(__A						):
         logger.error('''Vocabulary path ({}) should be a directory'''.format(__A						)						)
         return
     a					=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,)