




'''simple docstring'''


import inspect
from typing import List, Optional, Union

import torch
from torch import nn
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPImageProcessor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput

class 			A    (		nn.Module    ):







							def __init__(  self       :  str							,			lowerCAmelCase_       :  Dict							,			lowerCAmelCase_       :  str=1.0					)			->							Optional[int]:


														"""simple docstring"""


														super().__init__()

														_a	     =				cut_size
														_a	     =				cut_power







							def    __lowerCAmelCase			(  self       :  Union[str, Any]							,			lowerCAmelCase_       :  Optional[int]							,			lowerCAmelCase_       :  str					)			->							int:


														"""simple docstring"""


														_a     ,							_a	     =				pixel_values.shape[2:4]
														_a	     =				min(lowercase__							,			lowercase__					)
														_a	     =				min(lowercase__							,			lowercase__							,			self.cut_size					)
														_a	     =				[]
														for _ in range(lowercase__					):
																					_a	     =				int(torch.rand([]					) ** self.cut_power * (max_size - min_size) + min_size					)
																					_a	     =				torch.randint(0							,			sideX - size + 1							,			()					)
																					_a	     =				torch.randint(0							,			sideY - size + 1							,			()					)
																					_a	     =				pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]
																					cutouts.append(F.adaptive_avg_poolad(lowercase__							,			self.cut_size					)					)
														return torch.cat(lowercase__					)



def 		snake_case_    (UpperCamelCase       :  int ,						UpperCamelCase       :  Optional[int]      ):





							'''simple docstring'''


							_a	     =				F.normalize(SCREAMING_SNAKE_CASE_ ,						dim=-1      )
							_a	     =				F.normalize(SCREAMING_SNAKE_CASE_ ,						dim=-1      )
							return (x - y).norm(dim=-1      ).div(2      ).arcsin().pow(2      ).mul(2      )
def 		snake_case_    (UpperCamelCase       :  Dict ,						UpperCamelCase       :  int      ):





							'''simple docstring'''


							for param in model.parameters():
														_a	     =				value

class 			A    (		_a    ):







							def __init__(  self       :  List[str]							,			lowerCAmelCase_       :  Union[str, Any]							,			lowerCAmelCase_       :  Optional[int]							,			lowerCAmelCase_       :  Any							,			lowerCAmelCase_       :  Optional[int]							,			lowerCAmelCase_       :  int							,			lowerCAmelCase_       :  Optional[Any]							,			lowerCAmelCase_       :  List[str]							,			)			->							Optional[Any]:


														"""simple docstring"""


														super().__init__()
														self.register_modules(
														    vae=lowercase__							,			text_encoder=lowercase__							,			clip_model=lowercase__							,			tokenizer=lowercase__							,			unet=lowercase__							,			scheduler=lowercase__							,			feature_extractor=lowercase__							,			)

														_a	     =				transforms.Normalize(mean=feature_extractor.image_mean							,			std=feature_extractor.image_std					)
														_a	     =				(
														    feature_extractor.size
														    if isinstance(feature_extractor.size							,			lowercase__					)
														    else feature_extractor.size['''shortest_edge''']
														)
														_a	     =				MakeCutouts(self.cut_out_size					)

														set_requires_grad(self.text_encoder							,			lowercase__					)
														set_requires_grad(self.clip_model							,			lowercase__					)







							def    __lowerCAmelCase			(  self       :  Tuple							,			lowerCAmelCase_       :  List[Any] = "auto"					)			->							int:


														"""simple docstring"""


														if slice_size == "auto":
																					# half the attention head size is usually a good trade-off between
																					# speed and memory
																					_a	     =				self.unet.config.attention_head_dim // 2
														self.unet.set_attention_slice(lowercase__					)







							def    __lowerCAmelCase			(  self       :  int					)			->							Dict:


														"""simple docstring"""


														self.enable_attention_slicing(lowercase__					)







							def    __lowerCAmelCase			(  self       :  Optional[int]					)			->							Tuple:


														"""simple docstring"""


														set_requires_grad(self.vae							,			lowercase__					)







							def    __lowerCAmelCase			(  self       :  Optional[Any]					)			->							Any:


														"""simple docstring"""


														set_requires_grad(self.vae							,			lowercase__					)







							def    __lowerCAmelCase			(  self       :  Tuple					)			->							Tuple:


														"""simple docstring"""


														set_requires_grad(self.unet							,			lowercase__					)







							def    __lowerCAmelCase			(  self       :  Dict					)			->							List[Any]:


														"""simple docstring"""


														set_requires_grad(self.unet							,			lowercase__					)







							@torch.enable_grad()
							def    __lowerCAmelCase			(  self       :  Union[str, Any]							,			lowerCAmelCase_       :  int							,			lowerCAmelCase_       :  Tuple							,			lowerCAmelCase_       :  List[str]							,			lowerCAmelCase_       :  List[Any]							,			lowerCAmelCase_       :  Optional[int]							,			lowerCAmelCase_       :  Dict							,			lowerCAmelCase_       :  List[Any]							,			lowerCAmelCase_       :  int							,			lowerCAmelCase_       :  Optional[Any]=True							,			)			->							Optional[int]:


														"""simple docstring"""


														_a	     =				latents.detach().requires_grad_()

														_a	     =				self.scheduler.scale_model_input(lowercase__							,			lowercase__					)

														# predict the noise residual
														_a	     =				self.unet(lowercase__							,			lowercase__							,			encoder_hidden_states=lowercase__					).sample

														if isinstance(self.scheduler							,			(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)					):
																					_a	     =				self.scheduler.alphas_cumprod[timestep]
																					_a	     =				1 - alpha_prod_t
																					# compute predicted original sample from predicted noise also called
																					# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
																					_a	     =				(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

																					_a	     =				torch.sqrt(lowercase__					)
																					_a	     =				pred_original_sample * (fac) + latents * (1 - fac)
														elif isinstance(self.scheduler							,			lowercase__					):
																					_a	     =				self.scheduler.sigmas[index]
																					_a	     =				latents - sigma * noise_pred
														else:
																					raise ValueError(F'scheduler type {type(self.scheduler					)} not supported'					)

														_a	     =				1 / self.vae.config.scaling_factor * sample
														_a	     =				self.vae.decode(lowercase__					).sample
														_a	     =				(image / 2 + 0.5).clamp(0							,			1					)

														if use_cutouts:
																					_a	     =				self.make_cutouts(lowercase__							,			lowercase__					)
														else:
																					_a	     =				transforms.Resize(self.cut_out_size					)(lowercase__					)
														_a	     =				self.normalize(lowercase__					).to(latents.dtype					)

														_a	     =				self.clip_model.get_image_features(lowercase__					)
														_a	     =				image_embeddings_clip / image_embeddings_clip.norm(p=2							,			dim=-1							,			keepdim=lowercase__					)

														if use_cutouts:
																					_a	     =				spherical_dist_loss(lowercase__							,			lowercase__					)
																					_a	     =				dists.view([num_cutouts, sample.shape[0], -1]					)
																					_a	     =				dists.sum(2					).mean(0					).sum() * clip_guidance_scale
														else:
																					_a	     =				spherical_dist_loss(lowercase__							,			lowercase__					).mean() * clip_guidance_scale

														_a	     =				-torch.autograd.grad(lowercase__							,			lowercase__					)[0]

														if isinstance(self.scheduler							,			lowercase__					):
																					_a	     =				latents.detach() + grads * (sigma**2)
																					_a	     =				noise_pred_original
														else:
																					_a	     =				noise_pred_original - torch.sqrt(lowercase__					) * grads
														return noise_pred, latents







							@torch.no_grad()
							def __call__(  self       :  List[Any]							,			lowerCAmelCase_       :  Union[str, Any]							,			lowerCAmelCase_       :  int = 5_12							,			lowerCAmelCase_       :  Union[str, Any] = 5_12							,			lowerCAmelCase_       :  Any = 50							,			lowerCAmelCase_       :  Tuple = 7.5							,			lowerCAmelCase_       :  Dict = 1							,			lowerCAmelCase_       :  Optional[int] = 0.0							,			lowerCAmelCase_       :  Union[str, Any] = 1_00							,			lowerCAmelCase_       :  Union[str, Any] = None							,			lowerCAmelCase_       :  List[Any] = 4							,			lowerCAmelCase_       :  Union[str, Any] = True							,			lowerCAmelCase_       :  Optional[int] = None							,			lowerCAmelCase_       :  str = None							,			lowerCAmelCase_       :  Union[str, Any] = "pil"							,			lowerCAmelCase_       :  Any = True							,			)			->							Tuple:


														"""simple docstring"""


														if isinstance(lowercase__							,			lowercase__					):
																					_a	     =				1
														elif isinstance(lowercase__							,			lowercase__					):
																					_a	     =				len(lowercase__					)
														else:
																					raise ValueError(F'`prompt` has to be of type `str` or `list` but is {type(lowercase__					)}'					)

														if height % 8 != 0 or width % 8 != 0:
																					raise ValueError(F'`height` and `width` have to be divisible by 8 but are {height} and {width}.'					)

														# get prompt text embeddings
														_a	     =				self.tokenizer(
														    lowercase__							,			padding='''max_length'''							,			max_length=self.tokenizer.model_max_length							,			truncation=lowercase__							,			return_tensors='''pt'''							,			)
														_a	     =				self.text_encoder(text_input.input_ids.to(self.device					)					)[0]
														# duplicate text embeddings for each generation per prompt
														_a	     =				text_embeddings.repeat_interleave(lowercase__							,			dim=0					)

														if clip_guidance_scale > 0:
																					if clip_prompt is not None:
																												_a	     =				self.tokenizer(
																												    lowercase__							,			padding='''max_length'''							,			max_length=self.tokenizer.model_max_length							,			truncation=lowercase__							,			return_tensors='''pt'''							,			).input_ids.to(self.device					)
																					else:
																												_a	     =				text_input.input_ids.to(self.device					)
																					_a	     =				self.clip_model.get_text_features(lowercase__					)
																					_a	     =				text_embeddings_clip / text_embeddings_clip.norm(p=2							,			dim=-1							,			keepdim=lowercase__					)
																					# duplicate text embeddings clip for each generation per prompt
																					_a	     =				text_embeddings_clip.repeat_interleave(lowercase__							,			dim=0					)

														# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
														# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
														# corresponds to doing no classifier free guidance.
														_a	     =				guidance_scale > 1.0
														# get unconditional embeddings for classifier free guidance
														if do_classifier_free_guidance:
																					_a	     =				text_input.input_ids.shape[-1]
																					_a	     =				self.tokenizer(['''''']							,			padding='''max_length'''							,			max_length=lowercase__							,			return_tensors='''pt'''					)
																					_a	     =				self.text_encoder(uncond_input.input_ids.to(self.device					)					)[0]
																					# duplicate unconditional embeddings for each generation per prompt
																					_a	     =				uncond_embeddings.repeat_interleave(lowercase__							,			dim=0					)

																					# For classifier free guidance, we need to do two forward passes.
																					# Here we concatenate the unconditional and text embeddings into a single batch
																					# to avoid doing two forward passes
																					_a	     =				torch.cat([uncond_embeddings, text_embeddings]					)

														# get the initial random noise unless the user supplied it

														# Unlike in other pipelines, latents need to be generated in the target device
														# for 1-to-1 results reproducibility with the CompVis implementation.
														# However this currently doesn't work in `mps`.
														_a	     =				(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
														_a	     =				text_embeddings.dtype
														if latents is None:
																					if self.device.type == "mps":
																												# randn does not work reproducibly on mps
																												_a	     =				torch.randn(lowercase__							,			generator=lowercase__							,			device='''cpu'''							,			dtype=lowercase__					).to(
																												    self.device					)
																					else:
																												_a	     =				torch.randn(lowercase__							,			generator=lowercase__							,			device=self.device							,			dtype=lowercase__					)
														else:
																					if latents.shape != latents_shape:
																												raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'					)
																					_a	     =				latents.to(self.device					)

														# set timesteps
														_a	     =				'''offset''' in set(inspect.signature(self.scheduler.set_timesteps					).parameters.keys()					)
														_a	     =				{}
														if accepts_offset:
																					_a	     =				1

														self.scheduler.set_timesteps(lowercase__							,			**lowercase__					)

														# Some schedulers like PNDM have timesteps as arrays
														# It's more optimized to move all timesteps to correct device beforehand
														_a	     =				self.scheduler.timesteps.to(self.device					)

														# scale the initial noise by the standard deviation required by the scheduler
														_a	     =				latents * self.scheduler.init_noise_sigma

														# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
														# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
														# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
														# and should be between [0, 1]
														_a	     =				'''eta''' in set(inspect.signature(self.scheduler.step					).parameters.keys()					)
														_a	     =				{}
														if accepts_eta:
																					_a	     =				eta

														# check if the scheduler accepts generator
														_a	     =				'''generator''' in set(inspect.signature(self.scheduler.step					).parameters.keys()					)
														if accepts_generator:
																					_a	     =				generator

														for i, t in enumerate(self.progress_bar(lowercase__					)					):
																					# expand the latents if we are doing classifier free guidance
																					_a	     =				torch.cat([latents] * 2					) if do_classifier_free_guidance else latents
																					_a	     =				self.scheduler.scale_model_input(lowercase__							,			lowercase__					)

																					# predict the noise residual
																					_a	     =				self.unet(lowercase__							,			lowercase__							,			encoder_hidden_states=lowercase__					).sample

																					# perform classifier free guidance
																					if do_classifier_free_guidance:
																												_a     ,							_a	     =				noise_pred.chunk(2					)
																												_a	     =				noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																					# perform clip guidance
																					if clip_guidance_scale > 0:
																												_a	     =				(
																												    text_embeddings.chunk(2					)[1] if do_classifier_free_guidance else text_embeddings
																												)
																												_a     ,							_a	     =				self.cond_fn(
																												    lowercase__							,			lowercase__							,			lowercase__							,			lowercase__							,			lowercase__							,			lowercase__							,			lowercase__							,			lowercase__							,			lowercase__							,			)

																					# compute the previous noisy sample x_t -> x_t-1
																					_a	     =				self.scheduler.step(lowercase__							,			lowercase__							,			lowercase__							,			**lowercase__					).prev_sample

														# scale and decode the image latents with vae
														_a	     =				1 / self.vae.config.scaling_factor * latents
														_a	     =				self.vae.decode(lowercase__					).sample

														_a	     =				(image / 2 + 0.5).clamp(0							,			1					)
														_a	     =				image.cpu().permute(0							,			2							,			3							,			1					).numpy()

														if output_type == "pil":
																					_a	     =				self.numpy_to_pil(lowercase__					)

														if not return_dict:
																					return (image, None)

														return StableDiffusionPipelineOutput(images=lowercase__							,			nsfw_content_detected=lowercase__					)



"""simple docstring"""




import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    CpmAntConfig,
	    CpmAntForCausalLM,
	    CpmAntModel,
	    CpmAntTokenizer,
	)




@require_torch
class      UpperCamelCase	:

				"""simple docstring"""



				def __init__(				self		,	lowercase__		,	lowercase__=2		,	lowercase__=8		,	lowercase__=True		,	lowercase__=False		,	lowercase__=False		,	lowercase__=False		,	lowercase__=False		,	lowercase__=99		,	lowercase__=32		,	lowercase__=3		,	lowercase__=4		,	lowercase__=37		,	lowercase__=32		,	lowercase__=128		,	lowercase__=8		,	lowercase__=8		,	lowercase__=8		,	lowercase__=1.0		,	lowercase__=True		,	)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	parent
									SCREAMING_SNAKE_CASE				=	batch_size
									SCREAMING_SNAKE_CASE				=	seq_length
									SCREAMING_SNAKE_CASE				=	is_training
									SCREAMING_SNAKE_CASE				=	use_token_type_ids
									SCREAMING_SNAKE_CASE				=	use_input_mask
									SCREAMING_SNAKE_CASE				=	use_labels
									SCREAMING_SNAKE_CASE				=	use_mc_token_ids
									SCREAMING_SNAKE_CASE				=	vocab_size
									SCREAMING_SNAKE_CASE				=	hidden_size
									SCREAMING_SNAKE_CASE				=	num_hidden_layers
									SCREAMING_SNAKE_CASE				=	num_attention_heads
									SCREAMING_SNAKE_CASE				=	intermediate_size
									SCREAMING_SNAKE_CASE				=	num_buckets
									SCREAMING_SNAKE_CASE				=	max_distance
									SCREAMING_SNAKE_CASE				=	prompt_length
									SCREAMING_SNAKE_CASE				=	prompt_types
									SCREAMING_SNAKE_CASE				=	segment_types
									SCREAMING_SNAKE_CASE				=	init_std
									SCREAMING_SNAKE_CASE				=	return_dict


				def 							A					(				self			)    ->    Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	{}
									SCREAMING_SNAKE_CASE				=	ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size			).type(torch.intaa			)
									SCREAMING_SNAKE_CASE				=	False

									SCREAMING_SNAKE_CASE				=	self.get_config()

									return (config, input_ids)


				def 							A					(				self			)    ->    List[Any]:
									"""simple docstring"""







									return CpmAntConfig(
									    vocab_size=self.vocab_size		,	hidden_size=self.hidden_size		,	num_hidden_layers=self.num_hidden_layers		,	num_attention_heads=self.num_attention_heads		,	dim_ff=self.intermediate_size		,	position_bias_num_buckets=self.num_buckets		,	position_bias_max_distance=self.max_distance		,	prompt_types=self.prompt_types		,	prompt_length=self.prompt_length		,	segment_types=self.segment_types		,	use_cache=lowercase__		,	init_std=self.init_std		,	return_dict=self.return_dict		,	)


				def 							A					(				self		,	lowercase__		,	lowercase__		,	*lowercase__			)    ->    Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntModel(config=lowercase__			)
									model.to(lowercase__			)
									model.eval()

									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).last_hidden_state

									self.parent.assertEqual(hidden_states.shape		,	(self.batch_size, self.seq_length, config.hidden_size)			)


				def 							A					(				self		,	lowercase__		,	lowercase__		,	*lowercase__			)    ->    int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM(lowercase__			)
									model.to(lowercase__			)
									SCREAMING_SNAKE_CASE				=	input_ids['input_ids'].to(lowercase__			)
									model.eval()

									SCREAMING_SNAKE_CASE				=	model(**lowercase__			)
									self.parent.assertEqual(
									    model_output.logits.shape		,	(self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)		,	)




				def 							A					(				self			)    ->    Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.prepare_config_and_inputs()
									return config, inputs_dict






@require_torch
class      UpperCamelCase	(       __magic_name__  ,      __magic_name__  ,      unittest.TestCase      ):

				"""simple docstring"""



				UpperCAmelCase_				:      List[str]						  =				(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
				UpperCAmelCase_				:      List[Any]						  =				(
				    {"feature-extraction": CpmAntModel, "text-generation": CpmAntForCausalLM} if is_torch_available() else {}
				)

				UpperCAmelCase_				:      List[Any]						  =				False
				UpperCAmelCase_				:      Optional[int]						  =				False
				UpperCAmelCase_				:      Optional[Any]						  =				False
				UpperCAmelCase_				:      List[str]						  =				False
				UpperCAmelCase_				:      Dict						  =				False


				def 							A					(				self			)    ->    Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	CpmAntModelTester(self			)
									SCREAMING_SNAKE_CASE				=	ConfigTester(self		,	config_class=lowercase__			)


				def 							A					(				self			)    ->    Tuple:
									"""simple docstring"""







									self.config_tester.create_and_test_config_common_properties()
									self.config_tester.create_and_test_config_to_json_string()
									self.config_tester.create_and_test_config_to_json_file()
									self.config_tester.create_and_test_config_from_and_save_pretrained()
									self.config_tester.check_config_can_be_init_without_params()
									self.config_tester.check_config_arguments_init()


				def 							A					(				self			)    ->    List[str]:
									"""simple docstring"""







									unittest.skip('CPMAnt doesn\'t support input_embeds.'			)(self.test_inputs_embeds			)


				def 							A					(				self			)    ->    Optional[int]:
									"""simple docstring"""







									unittest.skip(
									    'CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.'			)(self.test_retain_grad_hidden_states_attentions			)


				def 							A					(				self			)    ->    List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_cpmant_model(lowercase__		,	lowercase__			)




				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE				=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_lm_head_model(lowercase__		,	lowercase__			)






@require_torch
class      UpperCamelCase	(       unittest.TestCase      ):

				"""simple docstring"""



				@tooslow
				def 							A					(				self			)    ->    Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'今天天气真好！'
									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntModel.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).last_hidden_state

									SCREAMING_SNAKE_CASE				=	torch.tensor(
									    [[[6.1_708, 5.9_244, 1.0_835], [6.5_207, 6.2_893, -11.3_324], [-1.0_107, -0.0_576, -5.9_577]]]		,	)
									self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,	lowercase__		,	atol=1E-2			)			)






@require_torch
class      UpperCamelCase	(       unittest.TestCase      ):

				"""simple docstring"""



				@tooslow
				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'今天天气真好！'
									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model(**lowercase__			).logits

									SCREAMING_SNAKE_CASE				=	torch.tensor(
									    [[[-6.4_267, -6.4_083, -6.3_958], [-5.8_802, -5.9_447, -5.7_811], [-5.3_896, -5.4_820, -5.4_295]]]		,	)
									self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,	lowercase__		,	atol=1E-2			)			)


				@tooslow
				def 							A					(				self			)    ->    Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	'今天天气不错，'
									SCREAMING_SNAKE_CASE				=	'今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'			)
									SCREAMING_SNAKE_CASE				=	model.generate(**lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer.batch_decode(lowercase__			)
									self.assertEqual(lowercase__		,	lowercase__			)




				@tooslow
				def 							A					(				self			)    ->    str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	'openbmb/cpm-ant-10b'
									SCREAMING_SNAKE_CASE				=	CpmAntForCausalLM.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	CpmAntTokenizer.from_pretrained(lowercase__			)
									SCREAMING_SNAKE_CASE				=	['今天天气不错，', '新年快乐，万事如意！']
									SCREAMING_SNAKE_CASE				=	[
									    '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的',
									    '新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的',
									]
									SCREAMING_SNAKE_CASE				=	tokenizer(lowercase__		,	return_tensors='pt'		,	padding=lowercase__			)
									SCREAMING_SNAKE_CASE				=	model.generate(**lowercase__			)
									SCREAMING_SNAKE_CASE				=	tokenizer.batch_decode(lowercase__			)
									self.assertEqual(lowercase__		,	lowercase__			)
