



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    __lowerCAmelCase					(	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 				snake_case  (				self		,							_snake_case=0.9999						):



									"""simple docstring"""





									_lowerCAmelCase         =       UNetaDConditionModel.from_pretrained(self.model_id		,							subfolder="""unet"""						)
									_lowerCAmelCase         =       unet.to(_snake_case						)
									_lowerCAmelCase         =       EMAModel(unet.parameters()		,							decay=_snake_case		,							model_cls=_snake_case		,							model_config=unet.config						)
									return unet, ema_unet




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       torch.randn(
									    self.batch_size		,							self.num_in_channels		,							self.latent_height		,							self.latent_width		,							generator=self.generator						).to(_snake_case						)
									_lowerCAmelCase         =       torch.randint(0		,							1000		,							size=(self.batch_size,)		,							generator=self.generator						).to(_snake_case						)
									_lowerCAmelCase         =       torch.randn(
									    self.batch_size		,							self.prompt_length		,							self.text_encoder_hidden_dim		,							generator=self.generator						).to(_snake_case						)
									return noisy_latents, timesteps, encoder_hidden_states




				def 				snake_case  (				self		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase         =       {}
									for k, param in unet.state_dict().items():
														_lowerCAmelCase         =       torch.randn_like(_snake_case						) + (param * torch.randn_like(_snake_case						))
														updated_state_dict.update({k: updated_param}						)
									unet.load_state_dict(_snake_case						)
									return unet




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       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 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       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()						)
									_lowerCAmelCase         =       list(unet.parameters()						)
									for s_param, param in zip(ema_unet.shadow_params		,							_snake_case						):
														assert torch.allclose(_snake_case		,							_snake_case						)

									# 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		,							_snake_case						):
														assert torch.allclose(_snake_case		,							_snake_case						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       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.
									_lowerCAmelCase         =       self.simulate_backprop(_snake_case						)

									# 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.
									_lowerCAmelCase         =       list(unet_pseudo_updated_step_one.parameters()						)
									for s_param, param in zip(ema_unet.shadow_params		,							_snake_case						):
														assert ~torch.allclose(_snake_case		,							_snake_case						)

									# 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		,							_snake_case						):
														assert ~torch.allclose(_snake_case		,							_snake_case						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       self.get_models()

									# First backprop + EMA
									_lowerCAmelCase         =       self.simulate_backprop(_snake_case						)
									ema_unet.step(unet_step_one.parameters()						)
									_lowerCAmelCase         =       ema_unet.shadow_params

									# Second backprop + EMA
									_lowerCAmelCase         =       self.simulate_backprop(_snake_case						)
									ema_unet.step(unet_step_two.parameters()						)
									_lowerCAmelCase         =       ema_unet.shadow_params

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




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       self.get_models(decay=0.0						)
									_lowerCAmelCase         =       self.simulate_backprop(_snake_case						)
									ema_unet.step(unet_step_one.parameters()						)
									_lowerCAmelCase         =       ema_unet.shadow_params

									_lowerCAmelCase         =       self.simulate_backprop(_snake_case						)
									ema_unet.step(unet_step_two.parameters()						)
									_lowerCAmelCase         =       ema_unet.shadow_params

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




				@skip_mps
				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       self.get_models()
									_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase         =       self.get_dummy_inputs()

									with tempfile.TemporaryDirectory() as tmpdir:
														ema_unet.save_pretrained(_snake_case						)
														_lowerCAmelCase         =       UNetaDConditionModel.from_pretrained(_snake_case		,							model_cls=_snake_case						)
														_lowerCAmelCase         =       loaded_unet.to(unet.device						)

									# Since no EMA step has been performed the outputs should match.
									_lowerCAmelCase         =       unet(_snake_case		,							_snake_case		,							_snake_case						).sample
									_lowerCAmelCase         =       loaded_unet(_snake_case		,							_snake_case		,							_snake_case						).sample

									assert torch.allclose(_snake_case		,							_snake_case		,							atol=1e-4						)








"""simple docstring"""







import os

from ...utils import logging
from .utils import DataProcessor, InputExample


A     =      logging.get_logger(__name__)






class   UpperCAmelCase__		(   UpperCamelCase  ):



		def __init__(			self	:     List[str]   ,				snake_case	:     List[Any]   ,				snake_case	:     str=None					)      ->    Tuple:






									'''simple docstring'''

									A    =       language
									A    =       train_language



		def        A_		(			self	:     Optional[int]   ,				snake_case	:     Any					)      ->    Union[str, Any]:






									'''simple docstring'''

									A    =       self.language if self.train_language is None else self.train_language
									A    =       self._read_tsv(os.path.join(snake_case   ,				f"""XNLI-MT-1.0/multinli/multinli.train.{lg}.tsv"""					)					)
									A    =       []
									for i, line in enumerate(snake_case					):
																if i == 0:
																							continue
																A    =       f"""train-{i}"""
																A    =       line[0]
																A    =       line[1]
																A    =       'contradiction' if line[2] == 'contradictory' else line[2]
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training input {text_a} is not a string"""					)
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training input {text_b} is not a string"""					)
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training label {label} is not a string"""					)
																examples.append(InputExample(guid=snake_case   ,				text_a=snake_case   ,				text_b=snake_case   ,				label=snake_case					)					)
									return examples



		def        A_		(			self	:     int   ,				snake_case	:     int					)      ->    Optional[int]:






									'''simple docstring'''

									A    =       self._read_tsv(os.path.join(snake_case   ,				'XNLI-1.0/xnli.test.tsv'					)					)
									A    =       []
									for i, line in enumerate(snake_case					):
																if i == 0:
																							continue
																A    =       line[0]
																if language != self.language:
																							continue
																A    =       f"""test-{i}"""
																A    =       line[6]
																A    =       line[7]
																A    =       line[1]
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training input {text_a} is not a string"""					)
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training input {text_b} is not a string"""					)
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training label {label} is not a string"""					)
																examples.append(InputExample(guid=snake_case   ,				text_a=snake_case   ,				text_b=snake_case   ,				label=snake_case					)					)
									return examples



		def        A_		(			self	:     Any					)      ->    List[Any]:






									'''simple docstring'''

									return ["contradiction", "entailment", "neutral"]



A     =      {
    'xnli': XnliProcessor,
}

A     =      {
    'xnli': 'classification',
}

A     =      {
    'xnli': 3,
}



