

"""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       ):
	a__		: List[str]					=					MODEL_FOR_CAUSAL_LM_MAPPING
	a__		: Optional[int]					=					TF_MODEL_FOR_CAUSAL_LM_MAPPING



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

		__UpperCAmelCase      		= text_generator(['''This is a test''', '''This is a second test''']						)
		self.assertEqual(
		    _lowercase			,    [
		        [
		            {
		                '''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      		= text_generator('''This is a test'''			,    do_sample=_lowercase			,    num_return_sequences=2			,    return_tensors=_lowercase						)
		self.assertEqual(
		    _lowercase			,    [
		        {'''generated_token_ids''': ANY(_lowercase						)},
		        {'''generated_token_ids''': ANY(_lowercase						)},
		    ]			,    )
		__UpperCAmelCase      		= text_generator.model.config.eos_token_id
		__UpperCAmelCase      		= '''<pad>'''
		__UpperCAmelCase      		= text_generator(
		    ['''This is a test''', '''This is a second test''']			,    do_sample=_lowercase			,    num_return_sequences=2			,    batch_size=2			,    return_tensors=_lowercase			,    )
		self.assertEqual(
		    _lowercase			,    [
		        [
		            {'''generated_token_ids''': ANY(_lowercase						)},
		            {'''generated_token_ids''': ANY(_lowercase						)},
		        ],
		        [
		            {'''generated_token_ids''': ANY(_lowercase						)},
		            {'''generated_token_ids''': ANY(_lowercase						)},
		        ],
		    ]			,    )



	@require_tf
	def        a  (  self     :       Dict						):
		__UpperCAmelCase      		= pipeline(task='''text-generation'''			,    model='''sshleifer/tiny-ctrl'''			,    framework='''tf'''						)

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

		__UpperCAmelCase      		= text_generator(['''This is a test''', '''This is a second test''']			,    do_sample=_lowercase						)
		self.assertEqual(
		    _lowercase			,    [
		        [
		            {
		                '''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			,    _lowercase     :       str			,    _lowercase     :       Optional[int]			,    _lowercase     :       str						):
		__UpperCAmelCase      		= TextGenerationPipeline(model=_lowercase			,    tokenizer=_lowercase						)
		return text_generator, ["This is a test", "Another test"]



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

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



	def        a  (  self     :       Optional[int]			,    _lowercase     :       List[Any]			,    _lowercase     :       str						):
		__UpperCAmelCase      		= text_generator.model
		__UpperCAmelCase      		= text_generator.tokenizer

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

		__UpperCAmelCase      		= text_generator('''This is a test'''			,    return_full_text=_lowercase						)
		self.assertEqual(_lowercase			,    [{'''generated_text''': ANY(_lowercase						)}]						)
		self.assertNotIn('''This is a test'''			,    outputs[0]['''generated_text''']						)

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

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

		__UpperCAmelCase      		= text_generator(['''This is great !''', '''Something else''']			,    num_return_sequences=2			,    do_sample=_lowercase						)
		self.assertEqual(
		    _lowercase			,    [
		        [{'''generated_text''': ANY(_lowercase						)}, {'''generated_text''': ANY(_lowercase						)}],
		        [{'''generated_text''': ANY(_lowercase						)}, {'''generated_text''': ANY(_lowercase						)}],
		    ]			,    )

		if text_generator.tokenizer.pad_token is not None:
			__UpperCAmelCase      		= text_generator(
			    ['''This is great !''', '''Something else''']			,    num_return_sequences=2			,    batch_size=2			,    do_sample=_lowercase						)
			self.assertEqual(
			    _lowercase			,    [
			        [{'''generated_text''': ANY(_lowercase						)}, {'''generated_text''': ANY(_lowercase						)}],
			        [{'''generated_text''': ANY(_lowercase						)}, {'''generated_text''': ANY(_lowercase						)}],
			    ]			,    )

		with self.assertRaises(_lowercase						):
			__UpperCAmelCase      		= text_generator('''test'''			,    return_full_text=_lowercase			,    return_text=_lowercase						)
		with self.assertRaises(_lowercase						):
			__UpperCAmelCase      		= text_generator('''test'''			,    return_full_text=_lowercase			,    return_tensors=_lowercase						)
		with self.assertRaises(_lowercase						):
			__UpperCAmelCase      		= text_generator('''test'''			,    return_text=_lowercase			,    return_tensors=_lowercase						)

		# 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      		= text_generator(''''''						)
			self.assertEqual(_lowercase			,    [{'''generated_text''': ANY(_lowercase						)}]						)
		else:
			with self.assertRaises((ValueError, AssertionError)						):
				__UpperCAmelCase      		= 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      		= ['''RwkvForCausalLM''', '''XGLMForCausalLM''', '''GPTNeoXForCausalLM''']
		if (
		    tokenizer.model_max_length < 1_00_00
		    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_00			,    max_new_tokens=20						)

			__UpperCAmelCase      		= text_generator('''This is a test''' * 5_00			,    handle_long_generation='''hole'''			,    max_new_tokens=20						)
			# Hole strategy cannot work
			with self.assertRaises(_lowercase						):
				text_generator(
				    '''This is a test''' * 5_00			,    handle_long_generation='''hole'''			,    max_new_tokens=tokenizer.model_max_length + 10			,    )



	@require_torch
	@require_accelerate
	@require_torch_gpu
	def        a  (  self     :       List[Any]						):
		import torch

		# Classic `model_kwargs`
		__UpperCAmelCase      		= 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      		= pipe('''This is a test'''						)
		self.assertEqual(
		    _lowercase			,    [
		        {
		            '''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      		= 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      		= pipe('''This is a test'''						)
		self.assertEqual(
		    _lowercase			,    [
		        {
		            '''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      		= 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      		= pipe('''This is a test'''						)
		self.assertEqual(
		    _lowercase			,    [
		        {
		            '''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     :       Union[str, Any]						):
		import torch

		__UpperCAmelCase      		= 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     :       int						):
		import torch

		__UpperCAmelCase      		= pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,    device_map='''auto'''			,    torch_dtype=torch.floataa						)
		pipe('''This is a test'''			,    do_sample=_lowercase			,    top_p=0.5						)



	def        a  (  self     :       int						):
		__UpperCAmelCase      		= '''Hello world'''
		__UpperCAmelCase      		= pipeline('''text-generation'''			,    model='''hf-internal-testing/tiny-random-gpt2'''						)
		if text_generator.model.framework == "tf":
			__UpperCAmelCase      		= logging.get_logger('''transformers.generation.tf_utils'''						)
		else:
			__UpperCAmelCase      		= logging.get_logger('''transformers.generation.utils'''						)
		__UpperCAmelCase      		= '''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(_lowercase						) as cl:
			__UpperCAmelCase      		= text_generator(_lowercase			,    max_length=10			,    max_new_tokens=1						)
		self.assertIn(_lowercase			,    cl.out						)

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

		with CaptureLogger(_lowercase						) as cl:
			__UpperCAmelCase      		= text_generator(_lowercase			,    max_length=10						)
		self.assertNotIn(_lowercase			,    cl.out						)



"""simple docstring"""





from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import ScoreSdeVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					lowercase__												=			42
					lowercase__												=			42







					def __init__(			self     :							int  ,lowercase_     :							UNetaDModel  ,lowercase_     :							ScoreSdeVeScheduler						):
						super().__init__()
						self.register_modules(unet=lowercase_  ,scheduler=lowercase_						)







					@torch.no_grad()
					def __call__(			self     :							List[str]  ,lowercase_     :							int = 1  ,lowercase_     :							int = 2_0_0_0  ,lowercase_     :							Optional[Union[torch.Generator, List[torch.Generator]]] = None  ,lowercase_     :							Optional[str] = "pil"  ,lowercase_     :							bool = True  ,**lowercase_     :							Dict  ,):
						lowerCAmelCase__      :		str    							=      self.unet.config.sample_size
						lowerCAmelCase__      :		int    							=      (batch_size, 3, img_size, img_size)

						lowerCAmelCase__      :		List[Any]    							=      self.unet

						lowerCAmelCase__      :		Tuple    							=      randn_tensor(lowercase_  ,generator=lowercase_						) * self.scheduler.init_noise_sigma
						lowerCAmelCase__      :		List[Any]    							=      sample.to(self.device						)

						self.scheduler.set_timesteps(lowercase_						)
						self.scheduler.set_sigmas(lowercase_						)

						for i, t in enumerate(self.progress_bar(self.scheduler.timesteps						)						):
							lowerCAmelCase__      :		Tuple    							=      self.scheduler.sigmas[i] * torch.ones(shape[0]  ,device=self.device						)

							# correction step
							for _ in range(self.scheduler.config.correct_steps						):
								lowerCAmelCase__      :		str    							=      self.unet(lowercase_  ,lowercase_						).sample
								lowerCAmelCase__      :		List[str]    							=      self.scheduler.step_correct(lowercase_  ,lowercase_  ,generator=lowercase_						).prev_sample

							# prediction step
							lowerCAmelCase__      :		Dict    							=      model(lowercase_  ,lowercase_						).sample
							lowerCAmelCase__      :		Optional[int]    							=      self.scheduler.step_pred(lowercase_  ,lowercase_  ,lowercase_  ,generator=lowercase_						)

							lowerCAmelCase__							,lowerCAmelCase__      :		Optional[Any]    							=      output.prev_sample, output.prev_sample_mean

						lowerCAmelCase__      :		List[str]    							=      sample_mean.clamp(0  ,1						)
						lowerCAmelCase__      :		Union[str, Any]    							=      sample.cpu().permute(0  ,2  ,3  ,1						).numpy()
						if output_type == "pil":
							lowerCAmelCase__      :		Union[str, Any]    							=      self.numpy_to_pil(lowercase_						)

						if not return_dict:
							return (sample,)

						return ImagePipelineOutput(images=lowercase_						)




