






'''simple docstring'''







import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class _lowercase      (							unittest.TestCase						):




							def 						lowerCamelCase_							(				self:							Union[str, Any]	):
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()




							@property
							def 						lowerCamelCase_							(				self:							int	):
														lowerCamelCase__   :      Optional[Any]			    =			1
														lowerCamelCase__   :      Any			    =			3
														lowerCamelCase__   :      str			    =			(32, 32)

														lowerCamelCase__   :      Union[str, Any]			    =			floats_tensor((batch_size, num_channels) + sizes     ,				rng=random.Random(0	)	).to(UpperCamelCase__	)
														return image




							@property
							def 						lowerCamelCase_							(				self:							Union[str, Any]	):
														torch.manual_seed(0	)
														lowerCamelCase__   :      str			    =			UNetaDConditionModel(
														    block_out_channels=(32, 64)     ,				layers_per_block=2     ,				sample_size=32     ,				in_channels=4     ,				out_channels=4     ,				down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")     ,				up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")     ,				cross_attention_dim=32     ,				)
														return model




							@property
							def 						lowerCamelCase_							(				self:							List[Any]	):
														torch.manual_seed(0	)
														lowerCamelCase__   :      Tuple			    =			AutoencoderKL(
														    block_out_channels=[32, 64]     ,				in_channels=3     ,				out_channels=3     ,				down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]     ,				up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]     ,				latent_channels=4     ,				)
														return model




							@property
							def 						lowerCamelCase_							(				self:							int	):
														torch.manual_seed(0	)
														lowerCamelCase__   :      Tuple			    =			CLIPTextConfig(
														    bos_token_id=0     ,				eos_token_id=2     ,				hidden_size=32     ,				intermediate_size=37     ,				layer_norm_eps=1e-05     ,				num_attention_heads=4     ,				num_hidden_layers=5     ,				pad_token_id=1     ,				vocab_size=1_000     ,				)
														return CLIPTextModel(UpperCamelCase__	)




							@property
							def 						lowerCamelCase_							(				self:							Optional[Any]	):
														def extract(*UpperCamelCase__:							int     ,				**UpperCamelCase__:							List[Any]	):

																					class _lowercase      :




																												def __init__(				self:							List[str]	):
																																			lowerCamelCase__   :      Tuple			    =			torch.ones([0]	)




																												def 						lowerCamelCase_							(				self:							List[str]     ,				UpperCamelCase__:							List[str]	):
																																			self.pixel_values.to(UpperCamelCase__	)
																																			return self



																					return Out()

														return extract




							def 						lowerCamelCase_							(				self:							Any	):
														lowerCamelCase__   :      List[str]			    =			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
														lowerCamelCase__   :      str			    =			self.dummy_cond_unet
														lowerCamelCase__   :      str			    =			DDIMScheduler(
														    beta_start=0.00_085     ,				beta_end=0.012     ,				beta_schedule="""scaled_linear"""     ,				clip_sample=UpperCamelCase__     ,				set_alpha_to_one=UpperCamelCase__     ,				)

														lowerCamelCase__   :      Optional[int]			    =			self.dummy_vae
														lowerCamelCase__   :      Optional[int]			    =			self.dummy_text_encoder
														lowerCamelCase__   :      Any			    =			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

														# make sure here that pndm scheduler skips prk
														lowerCamelCase__   :      Optional[Any]			    =			StableDiffusionPipeline(
														    unet=UpperCamelCase__     ,				scheduler=UpperCamelCase__     ,				vae=UpperCamelCase__     ,				text_encoder=UpperCamelCase__     ,				tokenizer=UpperCamelCase__     ,				safety_checker=UpperCamelCase__     ,				feature_extractor=self.dummy_extractor     ,				)
														lowerCamelCase__   :      int			    =			sd_pipe.to(UpperCamelCase__	)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      Optional[int]			    =			"""A painting of a squirrel eating a burger"""

														lowerCamelCase__   :      List[str]			    =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
														lowerCamelCase__   :      Tuple			    =			sd_pipe([prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=6.0     ,				num_inference_steps=2     ,				output_type="""np"""	)
														lowerCamelCase__   :      Dict			    =			output.images

														lowerCamelCase__   :      int			    =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
														lowerCamelCase__   :      List[str]			    =			sd_pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=6.0     ,				num_inference_steps=2     ,				output_type="""np"""     ,				return_dict=UpperCamelCase__     ,				)[0]

														lowerCamelCase__   :      List[str]			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      str			    =			image_from_tuple[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)
														lowerCamelCase__   :      Union[str, Any]			    =			np.array([0.5_753, 0.6_114, 0.5_001, 0.5_034, 0.5_470, 0.4_729, 0.4_971, 0.4_867, 0.4_867]	)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
														assert np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1e-2




							def 						lowerCamelCase_							(				self:							Union[str, Any]	):
														lowerCamelCase__   :      Dict			    =			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
														lowerCamelCase__   :      List[str]			    =			self.dummy_cond_unet
														lowerCamelCase__   :      Optional[int]			    =			PNDMScheduler(skip_prk_steps=UpperCamelCase__	)
														lowerCamelCase__   :      Optional[Any]			    =			self.dummy_vae
														lowerCamelCase__   :      Any			    =			self.dummy_text_encoder
														lowerCamelCase__   :      int			    =			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

														# make sure here that pndm scheduler skips prk
														lowerCamelCase__   :      Any			    =			StableDiffusionPipeline(
														    unet=UpperCamelCase__     ,				scheduler=UpperCamelCase__     ,				vae=UpperCamelCase__     ,				text_encoder=UpperCamelCase__     ,				tokenizer=UpperCamelCase__     ,				safety_checker=UpperCamelCase__     ,				feature_extractor=self.dummy_extractor     ,				)
														lowerCamelCase__   :      Optional[int]			    =			sd_pipe.to(UpperCamelCase__	)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      List[str]			    =			"""A painting of a squirrel eating a burger"""
														lowerCamelCase__   :      Optional[int]			    =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
														lowerCamelCase__   :      Tuple			    =			sd_pipe([prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=6.0     ,				num_inference_steps=2     ,				output_type="""np"""	)

														lowerCamelCase__   :      Dict			    =			output.images

														lowerCamelCase__   :      str			    =			torch.Generator(device=UpperCamelCase__	).manual_seed(0	)
														lowerCamelCase__   :      str			    =			sd_pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=6.0     ,				num_inference_steps=2     ,				output_type="""np"""     ,				return_dict=UpperCamelCase__     ,				)[0]

														lowerCamelCase__   :      List[Any]			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      List[str]			    =			image_from_tuple[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)
														lowerCamelCase__   :      str			    =			np.array([0.5_122, 0.5_712, 0.4_825, 0.5_053, 0.5_646, 0.4_769, 0.5_179, 0.4_894, 0.4_994]	)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
														assert np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1e-2




							def 						lowerCamelCase_							(				self:							List[Any]	):
														lowerCamelCase__   :      List[str]			    =			StableDiffusionPipeline.from_pretrained(
														    """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""     ,				safety_checker=UpperCamelCase__	)
														assert isinstance(UpperCamelCase__     ,				UpperCamelCase__	)
														assert isinstance(pipe.scheduler     ,				UpperCamelCase__	)
														assert pipe.safety_checker is None

														lowerCamelCase__   :      List[Any]			    =			pipe("""example prompt"""     ,				num_inference_steps=2	).images[0]
														assert image is not None

														# check that there's no error when saving a pipeline with one of the models being None
														with tempfile.TemporaryDirectory() as tmpdirname:
																					pipe.save_pretrained(UpperCamelCase__	)
																					lowerCamelCase__   :      Tuple			    =			StableDiffusionPipeline.from_pretrained(UpperCamelCase__	)

														# sanity check that the pipeline still works
														assert pipe.safety_checker is None
														lowerCamelCase__   :      Optional[Any]			    =			pipe("""example prompt"""     ,				num_inference_steps=2	).images[0]
														assert image is not None




							@unittest.skipIf(torch_device != """cuda"""     ,				"""This test requires a GPU"""	)
							def 						lowerCamelCase_							(				self:							Any	):
														lowerCamelCase__   :      Any			    =			self.dummy_cond_unet
														lowerCamelCase__   :      Any			    =			PNDMScheduler(skip_prk_steps=UpperCamelCase__	)
														lowerCamelCase__   :      Tuple			    =			self.dummy_vae
														lowerCamelCase__   :      List[Any]			    =			self.dummy_text_encoder
														lowerCamelCase__   :      str			    =			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

														# put models in fp16
														lowerCamelCase__   :      Optional[Any]			    =			unet.half()
														lowerCamelCase__   :      str			    =			vae.half()
														lowerCamelCase__   :      Any			    =			bert.half()

														# make sure here that pndm scheduler skips prk
														lowerCamelCase__   :      Any			    =			StableDiffusionPipeline(
														    unet=UpperCamelCase__     ,				scheduler=UpperCamelCase__     ,				vae=UpperCamelCase__     ,				text_encoder=UpperCamelCase__     ,				tokenizer=UpperCamelCase__     ,				safety_checker=UpperCamelCase__     ,				feature_extractor=self.dummy_extractor     ,				)
														lowerCamelCase__   :      List[str]			    =			sd_pipe.to(UpperCamelCase__	)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      Dict			    =			"""A painting of a squirrel eating a burger"""
														lowerCamelCase__   :      List[str]			    =			sd_pipe([prompt]     ,				num_inference_steps=2     ,				output_type="""np"""	).images

														assert image.shape == (1, 64, 64, 3)



@nightly
@require_torch_gpu
class _lowercase      (							unittest.TestCase						):




							def 						lowerCamelCase_							(				self:							List[str]	):
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()




							def 						lowerCamelCase_							(				self:							Any	):
														lowerCamelCase__   :      Dict			    =			"""cuda"""
														lowerCamelCase__   :      Dict			    =			StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
														lowerCamelCase__   :      Tuple			    =			pipe.to(UpperCamelCase__	)
														pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      Dict			    =			"""a photo of a cat"""
														lowerCamelCase__   :      str			    =			{
														    """editing_prompt""": ["""sunglasses"""],
														    """reverse_editing_direction""": [False],
														    """edit_warmup_steps""": 10,
														    """edit_guidance_scale""": 6,
														    """edit_threshold""": 0.95,
														    """edit_momentum_scale""": 0.5,
														    """edit_mom_beta""": 0.6,
														}

														lowerCamelCase__   :      str			    =			3
														lowerCamelCase__   :      List[str]			    =			7

														# no sega enabled
														lowerCamelCase__   :      List[str]			    =			torch.Generator(UpperCamelCase__	)
														generator.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      Optional[Any]			    =			pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=UpperCamelCase__     ,				num_inference_steps=50     ,				output_type="""np"""     ,				width=512     ,				height=512     ,				)

														lowerCamelCase__   :      Tuple			    =			output.images
														lowerCamelCase__   :      List[Any]			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      List[str]			    =			[
														    0.34_673_113,
														    0.38_492_733,
														    0.37_597_352,
														    0.34_086_335,
														    0.35_650_748,
														    0.35_579_205,
														    0.3_384_763,
														    0.34_340_236,
														    0.3_573_271,
														]

														assert image.shape == (1, 512, 512, 3)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2

														# with sega enabled
														# generator = torch.manual_seed(seed)
														generator.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      str			    =			pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=UpperCamelCase__     ,				num_inference_steps=50     ,				output_type="""np"""     ,				width=512     ,				height=512     ,				**UpperCamelCase__     ,				)

														lowerCamelCase__   :      Tuple			    =			output.images
														lowerCamelCase__   :      int			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      int			    =			[
														    0.41_887_826,
														    0.37_728_766,
														    0.30_138_272,
														    0.41_416_335,
														    0.41_664_985,
														    0.36_283_392,
														    0.36_191_246,
														    0.43_364_465,
														    0.43_001_732,
														]

														assert image.shape == (1, 512, 512, 3)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2




							def 						lowerCamelCase_							(				self:							Any	):
														lowerCamelCase__   :      Optional[Any]			    =			"""cuda"""
														lowerCamelCase__   :      Any			    =			StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
														lowerCamelCase__   :      str			    =			pipe.to(UpperCamelCase__	)
														pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      List[str]			    =			"""an image of a crowded boulevard, realistic, 4k"""
														lowerCamelCase__   :      int			    =			{
														    """editing_prompt""": """crowd, crowded, people""",
														    """reverse_editing_direction""": True,
														    """edit_warmup_steps""": 10,
														    """edit_guidance_scale""": 8.3,
														    """edit_threshold""": 0.9,
														    """edit_momentum_scale""": 0.5,
														    """edit_mom_beta""": 0.6,
														}

														lowerCamelCase__   :      Dict			    =			9
														lowerCamelCase__   :      str			    =			7

														# no sega enabled
														lowerCamelCase__   :      List[str]			    =			torch.Generator(UpperCamelCase__	)
														generator.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      List[str]			    =			pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=UpperCamelCase__     ,				num_inference_steps=50     ,				output_type="""np"""     ,				width=512     ,				height=512     ,				)

														lowerCamelCase__   :      List[Any]			    =			output.images
														lowerCamelCase__   :      Tuple			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      Dict			    =			[
														    0.43_497_998,
														    0.91_814_065,
														    0.7_540_739,
														    0.55_580_205,
														    0.8_467_265,
														    0.5_389_691,
														    0.62_574_506,
														    0.58_897_763,
														    0.50_926_757,
														]

														assert image.shape == (1, 512, 512, 3)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2

														# with sega enabled
														# generator = torch.manual_seed(seed)
														generator.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      str			    =			pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=UpperCamelCase__     ,				num_inference_steps=50     ,				output_type="""np"""     ,				width=512     ,				height=512     ,				**UpperCamelCase__     ,				)

														lowerCamelCase__   :      List[Any]			    =			output.images
														lowerCamelCase__   :      List[str]			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      str			    =			[
														    0.3_089_719,
														    0.30_500_144,
														    0.29_016_042,
														    0.30_630_964,
														    0.325_687,
														    0.29_419_225,
														    0.2_908_091,
														    0.28_723_598,
														    0.27_696_294,
														]

														assert image.shape == (1, 512, 512, 3)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2




							def 						lowerCamelCase_							(				self:							List[Any]	):
														lowerCamelCase__   :      List[Any]			    =			"""cuda"""
														lowerCamelCase__   :      Dict			    =			StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
														lowerCamelCase__   :      Union[str, Any]			    =			pipe.to(UpperCamelCase__	)
														pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      Optional[Any]			    =			"""a castle next to a river"""
														lowerCamelCase__   :      Union[str, Any]			    =			{
														    """editing_prompt""": ["""boat on a river, boat""", """monet, impression, sunrise"""],
														    """reverse_editing_direction""": False,
														    """edit_warmup_steps""": [15, 18],
														    """edit_guidance_scale""": 6,
														    """edit_threshold""": [0.9, 0.8],
														    """edit_momentum_scale""": 0.5,
														    """edit_mom_beta""": 0.6,
														}

														lowerCamelCase__   :      str			    =			48
														lowerCamelCase__   :      Optional[Any]			    =			7

														# no sega enabled
														lowerCamelCase__   :      int			    =			torch.Generator(UpperCamelCase__	)
														generator.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      Dict			    =			pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=UpperCamelCase__     ,				num_inference_steps=50     ,				output_type="""np"""     ,				width=512     ,				height=512     ,				)

														lowerCamelCase__   :      Any			    =			output.images
														lowerCamelCase__   :      Any			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      Tuple			    =			[
														    0.75_163_555,
														    0.76_037_145,
														    0.61_785,
														    0.9_189_673,
														    0.8_627_701,
														    0.85_189_694,
														    0.8_512_813,
														    0.87_012_076,
														    0.8_312_857,
														]

														assert image.shape == (1, 512, 512, 3)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2

														# with sega enabled
														# generator = torch.manual_seed(seed)
														generator.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      Optional[Any]			    =			pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=UpperCamelCase__     ,				num_inference_steps=50     ,				output_type="""np"""     ,				width=512     ,				height=512     ,				**UpperCamelCase__     ,				)

														lowerCamelCase__   :      Optional[Any]			    =			output.images
														lowerCamelCase__   :      Any			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      Dict			    =			[
														    0.73_553_365,
														    0.7_537_271,
														    0.74_341_905,
														    0.66_480_356,
														    0.6_472_925,
														    0.63_039_416,
														    0.64_812_905,
														    0.6_749_717,
														    0.6_517_102,
														]

														assert image.shape == (1, 512, 512, 3)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2




							def 						lowerCamelCase_							(				self:							int	):
														lowerCamelCase__   :      Optional[int]			    =			"""cuda"""
														lowerCamelCase__   :      Optional[int]			    =			StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""     ,				torch_dtype=torch.floataa	)
														lowerCamelCase__   :      int			    =			pipe.to(UpperCamelCase__	)
														pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      List[str]			    =			"""a photo of a cat"""
														lowerCamelCase__   :      int			    =			{
														    """editing_prompt""": ["""sunglasses"""],
														    """reverse_editing_direction""": [False],
														    """edit_warmup_steps""": 10,
														    """edit_guidance_scale""": 6,
														    """edit_threshold""": 0.95,
														    """edit_momentum_scale""": 0.5,
														    """edit_mom_beta""": 0.6,
														}

														lowerCamelCase__   :      int			    =			3
														lowerCamelCase__   :      str			    =			7

														# no sega enabled
														lowerCamelCase__   :      Tuple			    =			torch.Generator(UpperCamelCase__	)
														generator.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      List[Any]			    =			pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=UpperCamelCase__     ,				num_inference_steps=50     ,				output_type="""np"""     ,				width=512     ,				height=512     ,				)

														lowerCamelCase__   :      Optional[Any]			    =			output.images
														lowerCamelCase__   :      Tuple			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      Tuple			    =			[
														    0.34_887_695,
														    0.3_876_953,
														    0.375,
														    0.34_423_828,
														    0.3_581_543,
														    0.35_717_773,
														    0.3_383_789,
														    0.34_570_312,
														    0.359_375,
														]

														assert image.shape == (1, 512, 512, 3)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2

														# with sega enabled
														# generator = torch.manual_seed(seed)
														generator.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      Tuple			    =			pipe(
														    [prompt]     ,				generator=UpperCamelCase__     ,				guidance_scale=UpperCamelCase__     ,				num_inference_steps=50     ,				output_type="""np"""     ,				width=512     ,				height=512     ,				**UpperCamelCase__     ,				)

														lowerCamelCase__   :      Dict			    =			output.images
														lowerCamelCase__   :      int			    =			image[0, -3:, -3:, -1]
														lowerCamelCase__   :      Optional[Any]			    =			[
														    0.42_285_156,
														    0.36_914_062,
														    0.29_077_148,
														    0.42_041_016,
														    0.41_918_945,
														    0.35_498_047,
														    0.3_618_164,
														    0.4_423_828,
														    0.43_115_234,
														]

														assert image.shape == (1, 512, 512, 3)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2









'''simple docstring'''







import csv
import os

import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.csv import CsvDatasetReader, CsvDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase ) ->			Dict:
							assert isinstance(UpperCamelCase							,     UpperCamelCase )
							assert dataset.num_rows == 4
							assert dataset.num_columns == 3
							assert dataset.column_names == ["col_1", "col_2", "col_3"]
							for feature, expected_dtype in expected_features.items():
														assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize("""keep_in_memory"""							,     [False, True] )
def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Any:
							lowerCamelCase__   :      Any			    =			tmp_path / """cache"""
							lowerCamelCase__   :      int			    =			{"""col_1""": """int64""", """col_2""": """int64""", """col_3""": """float64"""}
							with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
														lowerCamelCase__   :      Any			    =			CsvDatasetReader(UpperCamelCase							,     cache_dir=UpperCamelCase							,     keep_in_memory=UpperCamelCase ).read()
							_check_csv_dataset(UpperCamelCase							,     UpperCamelCase )





@pytest.mark.parametrize(
    """features"""							,     [
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]							,     )
def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Any:
							lowerCamelCase__   :      Optional[int]			    =			tmp_path / """cache"""
							# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
							lowerCamelCase__   :      int			    =			{"""col_1""": """int64""", """col_2""": """int64""", """col_3""": """float64"""}
							lowerCamelCase__   :      int			    =			features.copy() if features else default_expected_features
							lowerCamelCase__   :      List[Any]			    =			(
							    Features({feature: Value(UpperCamelCase ) for feature, dtype in features.items()} ) if features is not None else None
							)
							lowerCamelCase__   :      Union[str, Any]			    =			CsvDatasetReader(UpperCamelCase							,     features=UpperCamelCase							,     cache_dir=UpperCamelCase ).read()
							_check_csv_dataset(UpperCamelCase							,     UpperCamelCase )





@pytest.mark.parametrize("""split"""							,     [None, NamedSplit("""train""" ), """train""", """test"""] )
def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Union[str, Any]:
							lowerCamelCase__   :      Any			    =			tmp_path / """cache"""
							lowerCamelCase__   :      Union[str, Any]			    =			{"""col_1""": """int64""", """col_2""": """int64""", """col_3""": """float64"""}
							lowerCamelCase__   :      List[Any]			    =			CsvDatasetReader(UpperCamelCase							,     cache_dir=UpperCamelCase							,     split=UpperCamelCase ).read()
							_check_csv_dataset(UpperCamelCase							,     UpperCamelCase )
							assert dataset.split == split if split else "train"





@pytest.mark.parametrize("""path_type"""							,     [str, list] )
def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Union[str, Any]:
							if issubclass(UpperCamelCase							,     UpperCamelCase ):
														lowerCamelCase__   :      int			    =			csv_path
							elif issubclass(UpperCamelCase							,     UpperCamelCase ):
														lowerCamelCase__   :      Tuple			    =			[csv_path]
							lowerCamelCase__   :      Optional[int]			    =			tmp_path / """cache"""
							lowerCamelCase__   :      str			    =			{"""col_1""": """int64""", """col_2""": """int64""", """col_3""": """float64"""}
							lowerCamelCase__   :      Dict			    =			CsvDatasetReader(UpperCamelCase							,     cache_dir=UpperCamelCase ).read()
							_check_csv_dataset(UpperCamelCase							,     UpperCamelCase )





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase=("train",) ) ->			int:
							assert isinstance(UpperCamelCase							,     UpperCamelCase )
							for split in splits:
														lowerCamelCase__   :      Optional[Any]			    =			dataset_dict[split]
														assert dataset.num_rows == 4
														assert dataset.num_columns == 3
														assert dataset.column_names == ["col_1", "col_2", "col_3"]
														for feature, expected_dtype in expected_features.items():
																					assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize("""keep_in_memory"""							,     [False, True] )
def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Optional[Any]:
							lowerCamelCase__   :      List[Any]			    =			tmp_path / """cache"""
							lowerCamelCase__   :      Dict			    =			{"""col_1""": """int64""", """col_2""": """int64""", """col_3""": """float64"""}
							with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
														lowerCamelCase__   :      Any			    =			CsvDatasetReader({"""train""": csv_path}							,     cache_dir=UpperCamelCase							,     keep_in_memory=UpperCamelCase ).read()
							_check_csv_datasetdict(UpperCamelCase							,     UpperCamelCase )





@pytest.mark.parametrize(
    """features"""							,     [
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]							,     )
def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			int:
							lowerCamelCase__   :      int			    =			tmp_path / """cache"""
							# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
							lowerCamelCase__   :      Optional[Any]			    =			{"""col_1""": """int64""", """col_2""": """int64""", """col_3""": """float64"""}
							lowerCamelCase__   :      List[str]			    =			features.copy() if features else default_expected_features
							lowerCamelCase__   :      Tuple			    =			(
							    Features({feature: Value(UpperCamelCase ) for feature, dtype in features.items()} ) if features is not None else None
							)
							lowerCamelCase__   :      str			    =			CsvDatasetReader({"""train""": csv_path}							,     features=UpperCamelCase							,     cache_dir=UpperCamelCase ).read()
							_check_csv_datasetdict(UpperCamelCase							,     UpperCamelCase )





@pytest.mark.parametrize("""split"""							,     [None, NamedSplit("""train""" ), """train""", """test"""] )
def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase							,     UpperCamelCase ) ->			Any:
							if split:
														lowerCamelCase__   :      Tuple			    =			{split: csv_path}
							else:
														lowerCamelCase__   :      Union[str, Any]			    =			{"""train""": csv_path, """test""": csv_path}
							lowerCamelCase__   :      Dict			    =			tmp_path / """cache"""
							lowerCamelCase__   :      Union[str, Any]			    =			{"""col_1""": """int64""", """col_2""": """int64""", """col_3""": """float64"""}
							lowerCamelCase__   :      Tuple			    =			CsvDatasetReader(UpperCamelCase							,     cache_dir=UpperCamelCase ).read()
							_check_csv_datasetdict(UpperCamelCase							,     UpperCamelCase							,     splits=list(path.keys() ) )
							assert all(dataset[split].split == split for split in path.keys() )





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase ) ->			Tuple:
							with open(UpperCamelCase							,     encoding="""utf-8""" ) as csvfile:
														yield from csv.reader(UpperCamelCase )





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase ) ->			Optional[int]:
							lowerCamelCase__   :      Union[str, Any]			    =			tmp_path / """cache"""
							lowerCamelCase__   :      List[str]			    =			os.path.join(UpperCamelCase							,     """tmp.csv""" )
							lowerCamelCase__   :      Optional[Any]			    =			CsvDatasetReader({"""train""": csv_path}							,     cache_dir=UpperCamelCase ).read()
							CsvDatasetWriter(dataset["""train"""]							,     UpperCamelCase							,     num_proc=1 ).write()

							lowerCamelCase__   :      Optional[Any]			    =			iter_csv_file(UpperCamelCase )
							lowerCamelCase__   :      str			    =			iter_csv_file(UpperCamelCase )

							for rowa, rowa in zip(UpperCamelCase							,     UpperCamelCase ):
														assert rowa == rowa





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase ) ->			List[str]:
							lowerCamelCase__   :      Any			    =			tmp_path / """cache"""
							lowerCamelCase__   :      List[str]			    =			os.path.join(UpperCamelCase							,     """tmp.csv""" )
							lowerCamelCase__   :      Optional[int]			    =			CsvDatasetReader({"""train""": csv_path}							,     cache_dir=UpperCamelCase ).read()
							CsvDatasetWriter(dataset["""train"""]							,     UpperCamelCase							,     num_proc=2 ).write()

							lowerCamelCase__   :      Tuple			    =			iter_csv_file(UpperCamelCase )
							lowerCamelCase__   :      Union[str, Any]			    =			iter_csv_file(UpperCamelCase )

							for rowa, rowa in zip(UpperCamelCase							,     UpperCamelCase ):
														assert rowa == rowa





def 					SCREAMING_SNAKE_CASE_			(UpperCamelCase							,     UpperCamelCase ) ->			Optional[int]:
							lowerCamelCase__   :      List[Any]			    =			tmp_path / """cache"""
							lowerCamelCase__   :      List[Any]			    =			os.path.join(UpperCamelCase							,     """tmp.csv""" )
							lowerCamelCase__   :      Dict			    =			CsvDatasetReader({"""train""": csv_path}							,     cache_dir=UpperCamelCase ).read()
							with pytest.raises(UpperCamelCase ):
														CsvDatasetWriter(dataset["""train"""]							,     UpperCamelCase							,     num_proc=0 )

