



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


			'''simple docstring'''

			_A      :		List[Any]			    =		MODEL_FOR_CAUSAL_LM_MAPPING
			_A      :		Dict			    =		TF_MODEL_FOR_CAUSAL_LM_MAPPING

			@require_torch
			def UpperCamelCase__					(							self     : Optional[Any]						):





					"""simple docstring"""






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





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		Optional[int]           =		pipeline(task="""text-generation"""							,     model="""sshleifer/tiny-ctrl"""							,     framework="""tf"""						)

					# Using `do_sample=False` to force deterministic output
					__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		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       :		Any           =		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 UpperCamelCase__					(							self     : List[Any]							,     lowerCAmelCase__     : Optional[int]							,     lowerCAmelCase__     : Tuple							,     lowerCAmelCase__     : Optional[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		Any           =		TextGenerationPipeline(model=lowerCAmelCase__							,     tokenizer=lowerCAmelCase__						)
					return text_generator, ["This is a test", "Another test"]

			def UpperCamelCase__					(							self     : str						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		Any           =		"""Hello I believe in"""
					__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		pipeline("""text-generation"""							,     model="""hf-internal-testing/tiny-random-gpt2"""						)
					__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		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       :		int           =		text_generator(lowerCAmelCase__							,     stop_sequence=""" fe"""						)
					self.assertEqual(lowerCAmelCase__							,     [{"""generated_text""": """Hello I believe in fe"""}]						)

			def UpperCamelCase__					(							self     : str							,     lowerCAmelCase__     : List[Any]							,     lowerCAmelCase__     : Optional[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		List[Any]           =		text_generator.model
					__SCREAMING_SNAKE_CASE       :		List[str]           =		text_generator.tokenizer

					__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		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       :		Optional[Any]           =		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       :		Union[str, Any]           =		pipeline(task="""text-generation"""							,     model=lowerCAmelCase__							,     tokenizer=lowerCAmelCase__							,     return_full_text=lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE       :		str           =		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       :		Dict           =		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       :		Tuple           =		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       :		Any           =		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       :		str           =		text_generator("""test"""							,     return_full_text=lowerCAmelCase__							,     return_text=lowerCAmelCase__						)
					with self.assertRaises(lowerCAmelCase__						):
							__SCREAMING_SNAKE_CASE       :		List[str]           =		text_generator("""test"""							,     return_full_text=lowerCAmelCase__							,     return_tensors=lowerCAmelCase__						)
					with self.assertRaises(lowerCAmelCase__						):
							__SCREAMING_SNAKE_CASE       :		Dict           =		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       :		List[Any]           =		text_generator(""""""						)
							self.assertEqual(lowerCAmelCase__							,     [{"""generated_text""": ANY(lowerCAmelCase__						)}]						)
					else:
							with self.assertRaises((ValueError, AssertionError)						):
									__SCREAMING_SNAKE_CASE       :		int           =		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       :		int           =		["""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						)

							__SCREAMING_SNAKE_CASE       :		Tuple           =		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(lowerCAmelCase__						):
									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 UpperCamelCase__					(							self     : str						):





					"""simple docstring"""






					import torch

					# Classic `model_kwargs`
					__SCREAMING_SNAKE_CASE       :		List[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						)
					__SCREAMING_SNAKE_CASE       :		str           =		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       :		Any           =		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       :		str           =		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       :		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						)
					__SCREAMING_SNAKE_CASE       :		Optional[int]           =		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 UpperCamelCase__					(							self     : Tuple						):





					"""simple docstring"""






					import torch

					__SCREAMING_SNAKE_CASE       :		str           =		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 UpperCamelCase__					(							self     : Optional[Any]						):





					"""simple docstring"""






					import torch

					__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		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 UpperCamelCase__					(							self     : List[str]						):





					"""simple docstring"""






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

					# The user only sets one -> no warning
					with CaptureLogger(lowerCAmelCase__						) as cl:
							__SCREAMING_SNAKE_CASE       :		List[Any]           =		text_generator(lowerCAmelCase__							,     max_new_tokens=1						)
					self.assertNotIn(lowerCAmelCase__							,     cl.out						)

					with CaptureLogger(lowerCAmelCase__						) as cl:
							__SCREAMING_SNAKE_CASE       :		int           =		text_generator(lowerCAmelCase__							,     max_length=1_0						)
					self.assertNotIn(lowerCAmelCase__							,     cl.out						)




'''simple docstring'''




import string


def 				lowerCAmelCase_   (	_lowerCamelCase:   str					):
		__SCREAMING_SNAKE_CASE       :		Dict           =		""""""
		for i in sequence:
				__SCREAMING_SNAKE_CASE       :		Any           =		ord(_lowerCamelCase					)
				if 65 <= extract <= 90:
						output += chr(1_55 - extract					)
				elif 97 <= extract <= 1_22:
						output += chr(2_19 - extract					)
				else:
						output += i
		return output


def 				lowerCAmelCase_   (	_lowerCamelCase:   str					):
		__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		string.ascii_letters
		__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		string.ascii_lowercase[::-1] + string.ascii_uppercase[::-1]
		return "".join(
		    letters_reversed[letters.index(_lowerCamelCase					)] if c in letters else c for c in sequence					)


def 				lowerCAmelCase_   (	):
		from timeit import timeit

		print("""Running performance benchmarks..."""					)
		__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		"""from string import printable ; from __main__ import atbash, atbash_slow"""
		print(F"> atbash_slow(): {timeit('atbash_slow(printable)'       ,      setup=_lowerCamelCase					)} seconds"					)
		print(F">      atbash(): {timeit('atbash(printable)'       ,      setup=_lowerCamelCase					)} seconds"					)


if __name__ == "__main__":
						for example in ("ABCDEFGH", "123GGjj", "testStringtest", "with space"):
												print(f"{example} encrypted in atbash: {atbash(example)}")
						benchmark()