



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


 """simple docstring"""


 UpperCAmelCase							=  MODEL_FOR_CAUSAL_LM_MAPPING
 UpperCAmelCase							=  TF_MODEL_FOR_CAUSAL_LM_MAPPING







 @require_torch
 def _snake_case (							self							)		->   str:
        _UpperCAmelCase  :						Optional[Any]			     = pipeline(task="""text-generation"""  ,model="""sshleifer/tiny-ctrl"""  ,framework="""pt"""							)
        # Using `do_sample=False` to force deterministic output
        _UpperCAmelCase  :						List[str]			     = 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  :						int			     = 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  :						Tuple			     = text_generator.model.config.eos_token_id
        _UpperCAmelCase  :						Union[str, Any]			     = """<pad>"""
        _UpperCAmelCase  :						Any			     = 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 _snake_case (							self							)		->   List[Any]:
        _UpperCAmelCase  :						Union[str, Any]			     = pipeline(task="""text-generation"""  ,model="""sshleifer/tiny-ctrl"""  ,framework="""tf"""							)

        # Using `do_sample=False` to force deterministic output
        _UpperCAmelCase  :						List[Any]			     = 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  :						Any			     = 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 _snake_case (							self  ,a_  ,a_  ,a_							)		->   Optional[int]:
        _UpperCAmelCase  :						Optional[Any]			     = TextGenerationPipeline(model=a_  ,tokenizer=a_							)
        return text_generator, ["This is a test", "Another test"]







 def _snake_case (							self							)		->   Tuple:
        _UpperCAmelCase  :						Dict			     = """Hello I believe in"""
        _UpperCAmelCase  :						Any			     = 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  :						Union[str, Any]			     = text_generator(a_  ,stop_sequence=""" fe"""							)
        self.assertEqual(a_  ,[{"""generated_text""": """Hello I believe in fe"""}]							)







 def _snake_case (							self  ,a_  ,a_							)		->   int:
        _UpperCAmelCase  :						Tuple			     = text_generator.model
        _UpperCAmelCase  :						Union[str, Any]			     = text_generator.tokenizer

        _UpperCAmelCase  :						Optional[int]			     = 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  :						Optional[Any]			     = 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  :						Optional[Any]			     = pipeline(task="""text-generation"""  ,model=a_  ,tokenizer=a_  ,return_full_text=a_							)
        _UpperCAmelCase  :						str			     = text_generator("""This is a test"""							)
        self.assertEqual(a_  ,[{"""generated_text""": ANY(a_							)}]							)
        self.assertNotIn("""This is a test"""  ,outputs[0]["""generated_text"""]							)

        _UpperCAmelCase  :						List[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  :						str			     = 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  :						str			     = text_generator("""test"""  ,return_full_text=a_  ,return_text=a_							)
        with self.assertRaises(a_							):
               _UpperCAmelCase  :						Dict			     = text_generator("""test"""  ,return_full_text=a_  ,return_tensors=a_							)
        with self.assertRaises(a_							):
               _UpperCAmelCase  :						Union[str, Any]			     = 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  :						Optional[Any]			     = text_generator(""""""							)
               self.assertEqual(a_  ,[{"""generated_text""": ANY(a_							)}]							)
        else:
               with self.assertRaises((ValueError, AssertionError)							):
                      _UpperCAmelCase  :						List[Any]			     = 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  :						int			     = ["""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							)

               _UpperCAmelCase  :						List[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 _snake_case (							self							)		->   Union[str, Any]:
        import torch

        # Classic `model_kwargs`
        _UpperCAmelCase  :						Tuple			     = 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  :						Union[str, 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"""
                    )
                }
            ]  ,)

        # Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
        _UpperCAmelCase  :						Dict			     = 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  :						Tuple			     = 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  :						int			     = 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  :						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"""
                    )
                }
            ]  ,)







 @require_torch
 @require_torch_gpu
 def _snake_case (							self							)		->   str:
        import torch

        _UpperCAmelCase  :						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 _snake_case (							self							)		->   Dict:
        import torch

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

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








'''simple docstring'''






import argparse

import torch

from transformers import RemBertConfig, RemBertModel, load_tf_weights_in_rembert
from transformers.utils import logging


logging.set_verbosity_info()




def     snake_case_					(     lowerCAmelCase_       ,   lowerCAmelCase_       ,   lowerCAmelCase_		)->		Union[str, Any]:





       '''simple docstring'''


       _UpperCAmelCase  :						int			     = RemBertConfig.from_json_file(lowerCAmelCase_		)
       print("""Building PyTorch model from configuration: {}""".format(str(lowerCAmelCase_		)		)		)
       _UpperCAmelCase  :						Any			     = RemBertModel(lowerCAmelCase_		)

       # Load weights from tf checkpoint
       load_tf_weights_in_rembert(lowerCAmelCase_       ,   lowerCAmelCase_       ,   lowerCAmelCase_		)

       # Save pytorch-model
       print("""Save PyTorch model to {}""".format(lowerCAmelCase_		)		)
       torch.save(model.state_dict()       ,   lowerCAmelCase_		)


if __name__ == "__main__":
   A_							: Optional[Any]						 =	argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
   )
   parser.add_argument(
       """--rembert_config_file""",
       default=None,
       type=str,
       required=True,
       help=(
           """The config json file corresponding to the pre-trained RemBERT model. \n"""
           """This specifies the model architecture."""
       ),
   )
   parser.add_argument(
       """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
   )
   A_							: Any						 =	parser.parse_args()
   convert_rembert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.rembert_config_file, args.pytorch_dump_path)



