

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




		UpperCAmelCase__      : List[str]			       =MODEL_FOR_CAUSAL_LM_MAPPING
		UpperCAmelCase__      : str			       =TF_MODEL_FOR_CAUSAL_LM_MAPPING


		@require_torch
		def  _lowercase      (			self    : List[str]		)					->Tuple:
								"""simple docstring"""







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

								SCREAMING_SNAKE_CASE :      List[str]       =   text_generator(["""This is a test""", """This is a second test"""]		)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        [
								            {
								                """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 :      List[Any]       =   text_generator("""This is a test"""				,				do_sample=UpperCAmelCase__				,				num_return_sequences=2				,				return_tensors=UpperCAmelCase__		)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        {"""generated_token_ids""": ANY(UpperCAmelCase__		)},
								        {"""generated_token_ids""": ANY(UpperCAmelCase__		)},
								    ]				,				)
								SCREAMING_SNAKE_CASE :      List[Any]       =   text_generator.model.config.eos_token_id
								SCREAMING_SNAKE_CASE :      List[Any]       =   """<pad>"""
								SCREAMING_SNAKE_CASE :      Dict       =   text_generator(
								    ["""This is a test""", """This is a second test"""]				,				do_sample=UpperCAmelCase__				,				num_return_sequences=2				,				batch_size=2				,				return_tensors=UpperCAmelCase__				,				)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        [
								            {"""generated_token_ids""": ANY(UpperCAmelCase__		)},
								            {"""generated_token_ids""": ANY(UpperCAmelCase__		)},
								        ],
								        [
								            {"""generated_token_ids""": ANY(UpperCAmelCase__		)},
								            {"""generated_token_ids""": ANY(UpperCAmelCase__		)},
								        ],
								    ]				,				)


		@require_tf
		def  _lowercase      (			self    : Dict		)					->Union[str, Any]:
								"""simple docstring"""







								SCREAMING_SNAKE_CASE :      List[Any]       =   pipeline(task="""text-generation"""				,				model="""sshleifer/tiny-ctrl"""				,				framework="""tf"""		)

								# Using `do_sample=False` to force deterministic output
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   text_generator("""This is a test"""				,				do_sample=UpperCAmelCase__		)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        {
								            """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=UpperCAmelCase__		)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        [
								            {
								                """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  _lowercase      (			self    : int				,				UpperCAmelCase__    : str				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : Any		)					->Optional[int]:
								"""simple docstring"""







								SCREAMING_SNAKE_CASE :      int       =   TextGenerationPipeline(model=UpperCAmelCase__				,				tokenizer=UpperCAmelCase__		)
								return text_generator, ["This is a test", "Another test"]


		def  _lowercase      (			self    : str		)					->int:
								"""simple docstring"""







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

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


		def  _lowercase      (			self    : Any				,				UpperCAmelCase__    : Tuple				,				UpperCAmelCase__    : List[Any]		)					->Any:
								"""simple docstring"""







								SCREAMING_SNAKE_CASE :      Union[str, Any]       =   text_generator.model
								SCREAMING_SNAKE_CASE :      str       =   text_generator.tokenizer

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

								SCREAMING_SNAKE_CASE :      Optional[int]       =   text_generator("""This is a test"""				,				return_full_text=UpperCAmelCase__		)
								self.assertEqual(UpperCAmelCase__				,				[{"""generated_text""": ANY(UpperCAmelCase__		)}]		)
								self.assertNotIn("""This is a test"""				,				outputs[0]["""generated_text"""]		)

								SCREAMING_SNAKE_CASE :      List[Any]       =   pipeline(task="""text-generation"""				,				model=UpperCAmelCase__				,				tokenizer=UpperCAmelCase__				,				return_full_text=UpperCAmelCase__		)
								SCREAMING_SNAKE_CASE :      Union[str, Any]       =   text_generator("""This is a test"""		)
								self.assertEqual(UpperCAmelCase__				,				[{"""generated_text""": ANY(UpperCAmelCase__		)}]		)
								self.assertNotIn("""This is a test"""				,				outputs[0]["""generated_text"""]		)

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

								SCREAMING_SNAKE_CASE :      List[str]       =   text_generator(["""This is great !""", """Something else"""]				,				num_return_sequences=2				,				do_sample=UpperCAmelCase__		)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        [{"""generated_text""": ANY(UpperCAmelCase__		)}, {"""generated_text""": ANY(UpperCAmelCase__		)}],
								        [{"""generated_text""": ANY(UpperCAmelCase__		)}, {"""generated_text""": ANY(UpperCAmelCase__		)}],
								    ]				,				)

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

								with self.assertRaises(UpperCAmelCase__		):
														SCREAMING_SNAKE_CASE :      Union[str, Any]       =   text_generator("""test"""				,				return_full_text=UpperCAmelCase__				,				return_text=UpperCAmelCase__		)
								with self.assertRaises(UpperCAmelCase__		):
														SCREAMING_SNAKE_CASE :      List[str]       =   text_generator("""test"""				,				return_full_text=UpperCAmelCase__				,				return_tensors=UpperCAmelCase__		)
								with self.assertRaises(UpperCAmelCase__		):
														SCREAMING_SNAKE_CASE :      Optional[int]       =   text_generator("""test"""				,				return_text=UpperCAmelCase__				,				return_tensors=UpperCAmelCase__		)

								# 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 :      Union[str, Any]       =   text_generator(""""""		)
														self.assertEqual(UpperCAmelCase__				,				[{"""generated_text""": ANY(UpperCAmelCase__		)}]		)
								else:
														with self.assertRaises((ValueError, AssertionError)		):
																				SCREAMING_SNAKE_CASE :      Optional[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 :      Any       =   ["""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 :      str       =   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(UpperCAmelCase__		):
																				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  _lowercase      (			self    : Optional[int]		)					->Tuple:
								"""simple docstring"""







								import torch

								# Classic `model_kwargs`
								SCREAMING_SNAKE_CASE :      Optional[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 :      Optional[int]       =   pipe("""This is a test"""		)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        {
								            """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 :      str       =   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 :      List[str]       =   pipe("""This is a test"""		)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        {
								            """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 :      List[str]       =   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 :      int       =   pipe("""This is a test"""		)
								self.assertEqual(
								    UpperCAmelCase__				,				[
								        {
								            """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  _lowercase      (			self    : int		)					->Optional[int]:
								"""simple docstring"""







								import torch

								SCREAMING_SNAKE_CASE :      List[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  _lowercase      (			self    : Any		)					->Dict:
								"""simple docstring"""







								import torch

								SCREAMING_SNAKE_CASE :      Optional[Any]       =   pipeline(model="""hf-internal-testing/tiny-random-bloom"""				,				device_map="""auto"""				,				torch_dtype=torch.floataa		)
								pipe("""This is a test"""				,				do_sample=UpperCAmelCase__				,				top_p=0.5		)




		def  _lowercase      (			self    : Any		)					->Dict:
								"""simple docstring"""







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

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

								with CaptureLogger(UpperCAmelCase__		) as cl:
														SCREAMING_SNAKE_CASE :      Optional[Any]       =   text_generator(UpperCAmelCase__				,				max_length=1_0		)
								self.assertNotIn(UpperCAmelCase__				,				cl.out		)







import os
import random
import sys

from . import cryptomath_module as cryptomath
from . import rabin_miller

UpperCAmelCase__ :     Optional[Any]				    =    3






def 					__lowercase      (     _A			)     ->						int:
      print("""Generating primitive root of p"""			)
      while True:
            SCREAMING_SNAKE_CASE :      int       =   random.randrange(3						,    _A			)
            if pow(_A						,    2						,    _A			) == 1:
                  continue
            if pow(_A						,    _A						,    _A			) == 1:
                  continue
            return g






def 					__lowercase      (     _A			)     ->						tuple[tuple[int, int, int, int], tuple[int, int]]:
      print("""Generating prime p..."""			)
      SCREAMING_SNAKE_CASE :      Optional[Any]       =   rabin_miller.generate_large_prime(_A			)  # select large prime number.
      SCREAMING_SNAKE_CASE :      Optional[Any]       =   primitive_root(_A			)  # one primitive root on modulo p.
      SCREAMING_SNAKE_CASE :      Optional[Any]       =   random.randrange(3						,    _A			)  # private_key -> have to be greater than 2 for safety.
      SCREAMING_SNAKE_CASE :      List[str]       =   cryptomath.find_mod_inverse(pow(_A						,    _A						,    _A			)						,    _A			)

      SCREAMING_SNAKE_CASE :      str       =   (key_size, e_a, e_a, p)
      SCREAMING_SNAKE_CASE :      Union[str, Any]       =   (key_size, d)

      return public_key, private_key






def 					__lowercase      (     _A						,    _A			)     ->						None:
      if os.path.exists(F"{name}_pubkey.txt"			) or os.path.exists(F"{name}_privkey.txt"			):
            print("""\nWARNING:"""			)
            print(
                F"\"{name}_pubkey.txt\" or \"{name}_privkey.txt\" already exists. \n"
                """Use a different name or delete these files and re-run this program."""			)
            sys.exit()

      SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      List[str]       =   generate_key(_A			)
      print(F"\nWriting public key to file {name}_pubkey.txt..."			)
      with open(F"{name}_pubkey.txt"						,    """w"""			) as fo:
            fo.write(F"{public_key[0]},{public_key[1]},{public_key[2]},{public_key[3]}"			)

      print(F"Writing private key to file {name}_privkey.txt..."			)
      with open(F"{name}_privkey.txt"						,    """w"""			) as fo:
            fo.write(F"{private_key[0]},{private_key[1]}"			)






def 					__lowercase      (     )     ->						None:
      print("""Making key files..."""			)
      make_key_files("""elgamal"""						,    2048			)
      print("""Key files generation successful"""			)


if __name__ == "__main__":
  main()




