


import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device


if is_torch_available():
			import torch

			from transformers import XLMProphetNetForConditionalGeneration, XLMProphetNetTokenizer



@require_torch
class      snake_case__	(   unittest.TestCase	):
				@slow
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				List[str]				=   XLMProphetNetForConditionalGeneration.from_pretrained('microsoft/xprophetnet-large-wiki100-cased'   )
										model.to(_UpperCAmelCase   )

										# encoder-decoder outputs
										__magic_name__					:				Any				=   torch.tensor([[17, 9_6208, 10_3471, 2]]   ).to(_UpperCAmelCase   )
										__magic_name__					:				Optional[Any]				=   torch.tensor(
										    [[2, 250, 9953, 34, 6_9489, 1620, 32, 11_8424, 624, 210, 105, 2913, 1032, 351]]   ).to(_UpperCAmelCase   )
										__magic_name__					:				List[str]				=   model(
										    input_ids=_UpperCAmelCase  ,   attention_mask=_UpperCAmelCase  ,   encoder_outputs=_UpperCAmelCase  ,   decoder_input_ids=_UpperCAmelCase   )
										__magic_name__					:				str				=   output[0]
										__magic_name__					:				Any				=   torch.Size((1, 14, 25_0012)   )
										self.assertEqual(output_predited_logis.shape  ,   _UpperCAmelCase   )
										__magic_name__					:				Union[str, Any]				=   torch.tensor(
										    [[[-6.3986, -8.2391, 12.5189], [-6.3289, -8.0864, 12.6211], [-6.2418, -8.0445, 12.7968]]]   ).to(_UpperCAmelCase   )
										self.assertTrue(torch.allclose(output_predited_logis[:, :3, :3]  ,   _UpperCAmelCase  ,   atol=1E-4   )   )

										# encoder outputs
										__magic_name__					:				Dict				=   model.prophetnet.encoder(_UpperCAmelCase   )[0]
										__magic_name__					:				Optional[Any]				=   torch.tensor(
										    [[[-1.4260, -0.7628, 0.8453], [-1.4719, -0.1391, 0.7807], [-1.7678, 0.0114, 0.4646]]]   ).to(_UpperCAmelCase   )
										__magic_name__					:				Dict				=   torch.Size((1, 4, 1024)   )
										self.assertEqual(encoder_outputs.shape  ,   _UpperCAmelCase   )
										self.assertTrue(torch.allclose(encoder_outputs[:, :3, :3]  ,   _UpperCAmelCase  ,   atol=1E-4   )   )

										# decoder outputs
										__magic_name__					:				Any				=   model.prophetnet.decoder(
										    _UpperCAmelCase  ,   encoder_hidden_states=_UpperCAmelCase  ,   )
										__magic_name__					:				List[Any]				=   decoder_outputs[1].view(1  ,   model.config.ngram  ,   14  ,   -1   )
										__magic_name__					:				List[str]				=   model.lm_head(_UpperCAmelCase   )
										__magic_name__					:				int				=   predicting_streams_logits[:, 0]
										self.assertTrue(torch.allclose(next_first_stream_logits[:, :3, :3]  ,   _UpperCAmelCase  ,   atol=1E-4   )   )
				@slow
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				Optional[int]				=   XLMProphetNetForConditionalGeneration.from_pretrained(
										    'microsoft/xprophetnet-large-wiki100-cased-xglue-ntg'   )
										model.to(_UpperCAmelCase   )

										__magic_name__					:				Dict				=   torch.tensor([[17, 9_6208, 10_3471, 2]]   ).to(_UpperCAmelCase   )
										__magic_name__					:				List[str]				=   torch.tensor(
										    [[2, 250, 9953, 34, 6_9489, 1620, 32, 11_8424, 624, 210, 105, 2913, 1032, 351]]   ).to(_UpperCAmelCase   )
										__magic_name__					:				Union[str, Any]				=   model(
										    input_ids=_UpperCAmelCase  ,   attention_mask=_UpperCAmelCase  ,   encoder_outputs=_UpperCAmelCase  ,   decoder_input_ids=_UpperCAmelCase   )
										__magic_name__					:				Tuple				=   output[0]
										__magic_name__					:				Dict				=   torch.Size((1, 14, 25_0012)   )
										self.assertEqual(output_predited_logis.shape  ,   _UpperCAmelCase   )
										# compare the actual values for a slice.
										__magic_name__					:				Dict				=   torch.tensor(
										    [[[-9.2253, -9.7173, -6.3529], [-7.6701, -9.0145, -1.9382], [-8.0195, -7.0004, -0.1523]]]   ).to(_UpperCAmelCase   )

										self.assertTrue(torch.allclose(output_predited_logis[:, :3, :3]  ,   _UpperCAmelCase  ,   atol=1E-4   )   )





				@slow
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				str				=   XLMProphetNetForConditionalGeneration.from_pretrained(
										    'microsoft/xprophetnet-large-wiki100-cased-xglue-ntg'   )
										model.to(_UpperCAmelCase   )
										__magic_name__					:				List[Any]				=   512

										__magic_name__					:				Tuple				=   XLMProphetNetTokenizer.from_pretrained('microsoft/xprophetnet-large-wiki100-cased-xglue-ntg'   )

										__magic_name__					:				int				=   (
										    'Microsoft Corporation intends to officially end free support for the Windows 7 operating system after'
										    ' January 14, 2020, according to the official portal of the organization. From that day, users of this'
										    ' system will not be able to receive security updates, which could make their computers vulnerable to'
										    ' cyber attacks.'
										)
										__magic_name__					:				Union[str, Any]				=   (
										    'орпорация Microsoft намерена официально прекратить бесплатную поддержку операционной системы Windows 7'
										    ' после 14 января 2020 года, сообщается на официальном портале организации . С указанного дня пользователи'
										    ' этой системы не смогут получать обновления безопасности, из-за чего их компьютеры могут стать уязвимыми'
										    ' к кибератакам.'
										)
										__magic_name__					:				List[str]				=   (
										    '根据该组织的官方门户网站，微软公司打算在2020年1月14日之后正式终止对Windows 7操作系统的免费支持。从那时起，该系统的用户将无法接收安全更新，这可能会使他们的计算机容易受到网络攻击。'
										)

										__magic_name__					:				Union[str, Any]				=   tokenizer(
										    [EN_SENTENCE, RU_SENTENCE, ZH_SENTENCE]  ,   padding=_UpperCAmelCase  ,   max_length=255  ,   return_tensors='pt'   ).input_ids
										__magic_name__					:				List[str]				=   input_ids.to(_UpperCAmelCase   )

										__magic_name__					:				str				=   model.generate(
										    _UpperCAmelCase  ,   num_beams=10  ,   length_penalty=1.0  ,   no_repeat_ngram_size=3  ,   early_stopping=_UpperCAmelCase   )
										__magic_name__					:				Tuple				=   [tokenizer.decode(_UpperCAmelCase  ,   skip_special_tokens=_UpperCAmelCase   ) for g in summary_ids]
										__magic_name__					:				str				=   'Microsoft to end Windows 7 free support after January 14, 2020'
										__magic_name__					:				Dict				=   'Microsoft намерена прекратить бесплатную поддержку Windows 7 после 14 января 2020 года'
										__magic_name__					:				Dict				=   '微软打算终止对Windows 7操作系统的免费支持'
										self.assertListEqual(
										    [EXPECTED_TITLE_EN, EXPECTED_TITLE_RU, EXPECTED_TITLE_ZH]  ,   _UpperCAmelCase  ,   )

										__magic_name__					:				Optional[Any]				=   model.generate(
										    _UpperCAmelCase  ,   num_beams=1  ,   length_penalty=1.0  ,   no_repeat_ngram_size=3  ,   early_stopping=_UpperCAmelCase   )
										__magic_name__					:				Optional[int]				=   [
										    tokenizer.convert_ids_to_tokens(_UpperCAmelCase  ,   skip_special_tokens=_UpperCAmelCase   ) for g in summary_ids_beama
										]
										__magic_name__					:				Optional[int]				=   '▁Microsoft ▁to ▁end ▁free ▁support ▁for ▁Windows ▁7'.split(' '   )
										__magic_name__					:				Optional[Any]				=   '▁Microsoft ▁намерен а ▁прекрати ть ▁бес плат ную ▁поддержку ▁Windows ▁7 ▁после ▁14 ▁января ▁2020 ▁года'.split(
										    ' '   )
										__magic_name__					:				Any				=   '微软 公司 打算 终止 对 Windows ▁7 操作 系统的 免费 支持'.split(' '   )
										self.assertListEqual(
										    [EXPECTED_TITLE_EN_BEAM1_TOK, EXPECTED_TITLE_RU_BEAM1_TOK, EXPECTED_TITLE_ZH_BEAM1_TOK]  ,   _UpperCAmelCase  ,   )






import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()



class      snake_case__	(   unittest.TestCase	):
				__lowerCAmelCase		=   '''hf-internal-testing/tiny-stable-diffusion-pipe'''
				__lowerCAmelCase		=   1
				__lowerCAmelCase		=   77
				__lowerCAmelCase		=   32
				__lowerCAmelCase		=   4
				__lowerCAmelCase		=   __lowerCAmelCase		=   64
				__lowerCAmelCase		=   torch.manual_seed(0	)
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase=0.9999   ):




										"""simple docstring"""




										__magic_name__					:				List[Any]				=   UNetaDConditionModel.from_pretrained(self.model_id  ,   subfolder='unet'   )
										__magic_name__					:				int				=   unet.to(_UpperCAmelCase   )
										__magic_name__					:				Optional[int]				=   EMAModel(unet.parameters()  ,   decay=_UpperCAmelCase  ,   model_cls=_UpperCAmelCase  ,   model_config=unet.config   )
										return unet, ema_unet
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				List[str]				=   torch.randn(
										    self.batch_size  ,   self.num_in_channels  ,   self.latent_height  ,   self.latent_width  ,   generator=self.generator   ).to(_UpperCAmelCase   )
										__magic_name__					:				List[str]				=   torch.randint(0  ,   1000  ,   size=(self.batch_size,)  ,   generator=self.generator   ).to(_UpperCAmelCase   )
										__magic_name__					:				Any				=   torch.randn(
										    self.batch_size  ,   self.prompt_length  ,   self.text_encoder_hidden_dim  ,   generator=self.generator   ).to(_UpperCAmelCase   )
										return noisy_latents, timesteps, encoder_hidden_states
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				Any				=   {}
										for k, param in unet.state_dict().items():
																__magic_name__					:				str				=   torch.randn_like(_UpperCAmelCase   ) + (param * torch.randn_like(_UpperCAmelCase   ))
																updated_state_dict.update({k: updated_param}   )
										unet.load_state_dict(_UpperCAmelCase   )
										return unet
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				int				=   self.get_models()
										# Take the first (hypothetical) EMA step.
										ema_unet.step(unet.parameters()   )
										assert ema_unet.optimization_step == 1

										# Take two more.
										for _ in range(2   ):
																ema_unet.step(unet.parameters()   )
										assert ema_unet.optimization_step == 3
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Union[str, Any]				=   self.get_models()
										# Since the `unet` is not being updated (i.e., backprop'd)
										# there won't be any difference between the `params` of `unet`
										# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
										ema_unet.step(unet.parameters()   )
										__magic_name__					:				Any				=   list(unet.parameters()   )
										for s_param, param in zip(ema_unet.shadow_params  ,   _UpperCAmelCase   ):
																assert torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )

										# The above holds true even if we call `ema.step()` multiple times since
										# `unet` params are still not being updated.
										for _ in range(4   ):
																ema_unet.step(unet.parameters()   )
										for s_param, param in zip(ema_unet.shadow_params  ,   _UpperCAmelCase   ):
																assert torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[Any]				=   self.get_models()
										# Here we simulate the parameter updates for `unet`. Since there might
										# be some parameters which are initialized to zero we take extra care to
										# initialize their values to something non-zero before the multiplication.
										__magic_name__					:				Any				=   self.simulate_backprop(_UpperCAmelCase   )

										# Take the EMA step.
										ema_unet.step(unet_pseudo_updated_step_one.parameters()   )

										# Now the EMA'd parameters won't be equal to the original model parameters.
										__magic_name__					:				Any				=   list(unet_pseudo_updated_step_one.parameters()   )
										for s_param, param in zip(ema_unet.shadow_params  ,   _UpperCAmelCase   ):
																assert ~torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )

										# Ensure this is the case when we take multiple EMA steps.
										for _ in range(4   ):
																ema_unet.step(unet.parameters()   )
										for s_param, param in zip(ema_unet.shadow_params  ,   _UpperCAmelCase   ):
																assert ~torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				List[str]				=   self.get_models()

										# First backprop + EMA
										__magic_name__					:				int				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_one.parameters()   )
										__magic_name__					:				Any				=   ema_unet.shadow_params

										# Second backprop + EMA
										__magic_name__					:				Any				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_two.parameters()   )
										__magic_name__					:				Union[str, Any]				=   ema_unet.shadow_params

										for step_one, step_two in zip(_UpperCAmelCase  ,   _UpperCAmelCase   ):
																assert ~torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[Any]				=   self.get_models(decay=0.0   )
										__magic_name__					:				str				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_one.parameters()   )
										__magic_name__					:				List[Any]				=   ema_unet.shadow_params

										__magic_name__					:				Optional[Any]				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_two.parameters()   )
										__magic_name__					:				Optional[Any]				=   ema_unet.shadow_params

										for step_one, step_two in zip(_UpperCAmelCase  ,   _UpperCAmelCase   ):
																assert torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )





				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[int]				=   self.get_models()
										__magic_name__		,			__magic_name__		,			__magic_name__					:				str				=   self.get_dummy_inputs()

										with tempfile.TemporaryDirectory() as tmpdir:
																ema_unet.save_pretrained(_UpperCAmelCase   )
																__magic_name__					:				Optional[int]				=   UNetaDConditionModel.from_pretrained(_UpperCAmelCase  ,   model_cls=_UpperCAmelCase   )
																__magic_name__					:				str				=   loaded_unet.to(unet.device   )

										# Since no EMA step has been performed the outputs should match.
										__magic_name__					:				Union[str, Any]				=   unet(_UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ).sample
										__magic_name__					:				List[Any]				=   loaded_unet(_UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ).sample

										assert torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase  ,   atol=1E-4   )


