





import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


_SCREAMING_SNAKE_CASE      	=		datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
      import pyspark
@dataclass
class SCREAMING_SNAKE_CASE_						(					datasets.BuilderConfig					):
      __magic_name__:			Optional[datasets.Features]       =       None


def 							SCREAMING_SNAKE_CASE__						(    __a   , __a   , ):
 import pyspark

 def generate_fn():
  snake_case_		:			Union[str, Any]											=	df.select('*'   , pyspark.sql.functions.spark_partition_id().alias('part_id'   )   )
  for partition_id in partition_order:
   snake_case_		:			Tuple											=	df_with_partition_id.select('*'   ).where(f"""part_id = {partition_id}"""   ).drop('part_id'   )
   snake_case_		:			int											=	partition_df.collect()
   snake_case_		:			int											=	0
   for row in rows:
    yield f"""{partition_id}_{row_id}""", row.asDict()
    row_id += 1

 return generate_fn
class SCREAMING_SNAKE_CASE_						(					_BaseExamplesIterable					):






      def __init__(      self      :      Union[str, Any]		,   _A      :      "pyspark.sql.DataFrame"		,   _A      :      Optional[Any]=None		,   )							->    Optional[int]:







       """simple docstring"""


       snake_case_		:			Any											=	df
       snake_case_		:			Any											=	partition_order or range(self.df.rdd.getNumPartitions()	)
       snake_case_		:			int											=	_generate_iterable_examples(self.df		,   self.partition_order	)






      def __iter__(      self      :      List[Any]	)							->    int:







       """simple docstring"""


       yield from self.generate_examples_fn()






      def UpperCAmelCase_				(      self      :      List[str]		,   _A      :      np.random.Generator	)							->    "SparkExamplesIterable":







       """simple docstring"""


       snake_case_		:			Union[str, Any]											=	list(range(self.df.rdd.getNumPartitions()	)	)
       generator.shuffle(_A	)
       return SparkExamplesIterable(self.df		,   partition_order=_A	)






      def UpperCAmelCase_				(      self      :      Optional[int]		,   _A      :      int		,   _A      :      int	)							->    "SparkExamplesIterable":







       """simple docstring"""


       snake_case_		:			Tuple											=	self.split_shard_indices_by_worker(_A		,   _A	)
       return SparkExamplesIterable(self.df		,   partition_order=_A	)






      @property
      def UpperCAmelCase_				(      self      :      List[Any]	)							->    int:







       """simple docstring"""


       return len(self.partition_order	)
class SCREAMING_SNAKE_CASE_						(					datasets.DatasetBuilder					):
      __magic_name__:			Dict       =       SparkConfig






      def __init__(      self      :      int		,   _A      :      "pyspark.sql.DataFrame"		,   _A      :      str = None		,   _A      :      str = None		,   **_A      :      Dict		,   )							->    Union[str, Any]:







       """simple docstring"""


       import pyspark

       snake_case_		:			List[str]											=	pyspark.sql.SparkSession.builder.getOrCreate()
       snake_case_		:			List[Any]											=	df
       snake_case_		:			Tuple											=	working_dir

       super().__init__(
           cache_dir=_A		,   config_name=str(self.df.semanticHash()	)		,   **_A		,   )






      def UpperCAmelCase_				(      self      :      List[str]	)							->    Optional[Any]:







       """simple docstring"""


       def create_cache_and_write_probe(_A      :      Optional[Any]	):
        # makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
        # already exist.
        os.makedirs(self._cache_dir		,   exist_ok=_A	)
        snake_case_		:			Any											=	os.path.join(self._cache_dir		,   'fs_test' + uuid.uuida().hex	)
        # Opening the file in append mode will create a new file unless it already exists, in which case it will not
        # change the file contents.
        open(_A		,   'a'	)
        return [probe_file]

       if self._spark.conf.get('spark.master'		,   ''	).startswith('local'	):
        return

       # If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
       # accessible to the driver.
       # TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
       if self._cache_dir:
        snake_case_		:			Any											=	(
            self._spark.sparkContext.parallelize(range(1	)		,   1	).mapPartitions(_A	).collect()
        )
        if os.path.isfile(probe[0]	):
         return

       raise ValueError(
           'When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir'	)






      def UpperCAmelCase_				(      self      :      Union[str, Any]	)							->    List[str]:







       """simple docstring"""


       return datasets.DatasetInfo(features=self.config.features	)






      def UpperCAmelCase_				(      self      :      List[Any]		,   _A      :      datasets.download.download_manager.DownloadManager	)							->    Tuple:







       """simple docstring"""


       return [datasets.SplitGenerator(name=datasets.Split.TRAIN	)]






      def UpperCAmelCase_				(      self      :      Any		,   _A      :      Dict	)							->    Tuple:







       """simple docstring"""


       import pyspark

       def get_arrow_batch_size(_A      :      List[str]	):
        for batch in it:
         yield pa.RecordBatch.from_pydict({'batch_bytes': [batch.nbytes]}	)

       snake_case_		:			Union[str, Any]											=	self.df.count()
       snake_case_		:			Optional[Any]											=	df_num_rows if df_num_rows <= 100 else 100
       # Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
       snake_case_		:			int											=	(
           self.df.limit(_A	)
           .repartition(1	)
           .mapInArrow(_A		,   'batch_bytes: long'	)
           .agg(pyspark.sql.functions.sum('batch_bytes'	).alias('sample_bytes'	)	)
           .collect()[0]
           .sample_bytes
           / sample_num_rows
       )
       snake_case_		:			List[Any]											=	approx_bytes_per_row * df_num_rows
       if approx_total_size > max_shard_size:
        # Make sure there is at least one row per partition.
        snake_case_		:			int											=	min(_A		,   int(approx_total_size / max_shard_size	)	)
        snake_case_		:			str											=	self.df.repartition(_A	)






      def UpperCAmelCase_				(      self      :      List[str]		,   _A      :      str		,   _A      :      str		,   _A      :      int		,   )							->    Iterable[Tuple[int, bool, Union[int, tuple]]]:







       """simple docstring"""


       import pyspark

       snake_case_		:			Optional[Any]											=	ParquetWriter if file_format == 'parquet' else ArrowWriter
       snake_case_		:			Dict											=	os.path.join(self._working_dir		,   os.path.basename(_A	)	) if self._working_dir else fpath
       snake_case_		:			List[str]											=	file_format == 'parquet'

       # Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
       # pickling the SparkContext.
       snake_case_		:			Optional[Any]											=	self.config.features
       snake_case_		:			List[str]											=	self._writer_batch_size
       snake_case_		:			Dict											=	self._fs.storage_options

       def write_arrow(_A      :      Optional[int]	):
        # Within the same SparkContext, no two task attempts will share the same attempt ID.
        snake_case_		:			Dict											=	pyspark.TaskContext().taskAttemptId()
        snake_case_		:			Tuple											=	next(_A		,   _A	)
        if first_batch is None:
         # Some partitions might not receive any data.
         return pa.RecordBatch.from_arrays(
             [[task_id], [0], [0]]		,   names=['task_id', 'num_examples', 'num_bytes']		,   )
        snake_case_		:			Optional[Any]											=	0
        snake_case_		:			Dict											=	writer_class(
            features=_A		,   path=working_fpath.replace('SSSSS'		,   F"""{shard_id:05d}"""	).replace('TTTTT'		,   F"""{task_id:05d}"""	)		,   writer_batch_size=_A		,   storage_options=_A		,   embed_local_files=_A		,   )
        snake_case_		:			List[Any]											=	pa.Table.from_batches([first_batch]	)
        writer.write_table(_A	)
        for batch in it:
         if max_shard_size is not None and writer._num_bytes >= max_shard_size:
          snake_case_		:			List[str]											=	writer.finalize()
          writer.close()
          yield pa.RecordBatch.from_arrays(
              [[task_id], [num_examples], [num_bytes]]		,   names=['task_id', 'num_examples', 'num_bytes']		,   )
          shard_id += 1
          snake_case_		:			Tuple											=	writer_class(
              features=writer._features		,   path=working_fpath.replace('SSSSS'		,   F"""{shard_id:05d}"""	).replace('TTTTT'		,   F"""{task_id:05d}"""	)		,   writer_batch_size=_A		,   storage_options=_A		,   embed_local_files=_A		,   )
         snake_case_		:			List[Any]											=	pa.Table.from_batches([batch]	)
         writer.write_table(_A	)

        if writer._num_bytes > 0:
         snake_case_		:			str											=	writer.finalize()
         writer.close()
         yield pa.RecordBatch.from_arrays(
             [[task_id], [num_examples], [num_bytes]]		,   names=['task_id', 'num_examples', 'num_bytes']		,   )

        if working_fpath != fpath:
         for file in os.listdir(os.path.dirname(_A	)	):
          snake_case_		:			str											=	os.path.join(os.path.dirname(_A	)		,   os.path.basename(_A	)	)
          shutil.move(_A		,   _A	)

       snake_case_		:			int											=	(
           self.df.mapInArrow(_A		,   'task_id: long, num_examples: long, num_bytes: long'	)
           .groupBy('task_id'	)
           .agg(
               pyspark.sql.functions.sum('num_examples'	).alias('total_num_examples'	)		,   pyspark.sql.functions.sum('num_bytes'	).alias('total_num_bytes'	)		,   pyspark.sql.functions.count('num_bytes'	).alias('num_shards'	)		,   pyspark.sql.functions.collect_list('num_examples'	).alias('shard_lengths'	)		,   )
           .collect()
       )
       for row in stats:
        yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)






      def UpperCAmelCase_				(      self      :      Union[str, Any]		,   _A      :      "datasets.SplitGenerator"		,   _A      :      str = "arrow"		,   _A      :      Optional[Union[str, int]] = None		,   _A      :      Optional[int] = None		,   **_A      :      Optional[int]		,   )							->    Dict:







       """simple docstring"""


       self._validate_cache_dir()

       snake_case_		:			List[Any]											=	convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE	)
       self._repartition_df_if_needed(_A	)
       snake_case_		:			Tuple											=	not is_remote_filesystem(self._fs	)
       snake_case_		:			str											=	os.path.join if is_local else posixpath.join

       snake_case_		:			Optional[int]											=	'-TTTTT-SSSSS-of-NNNNN'
       snake_case_		:			List[str]											=	F"""{self.name}-{split_generator.name}{SUFFIX}.{file_format}"""
       snake_case_		:			str											=	path_join(self._output_dir		,   _A	)

       snake_case_		:			Any											=	0
       snake_case_		:			List[str]											=	0
       snake_case_		:			Tuple											=	0
       snake_case_		:			Tuple											=	[]
       snake_case_		:			Tuple											=	[]

       for task_id, content in self._prepare_split_single(_A		,   _A		,   _A	):
        (
            snake_case_
        )		:			int											=	content
        if num_bytes > 0:
         total_num_examples += num_examples
         total_num_bytes += num_bytes
         total_shards += num_shards
         task_id_and_num_shards.append((task_id, num_shards)	)
         all_shard_lengths.extend(_A	)

       snake_case_		:			List[str]											=	total_num_examples
       snake_case_		:			Union[str, Any]											=	total_num_bytes

       # should rename everything at the end
       logger.debug(F"""Renaming {total_shards} shards."""	)
       if total_shards > 1:
        snake_case_		:			Dict											=	all_shard_lengths

        # Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
        # pickling error due to pickling the SparkContext.
        snake_case_		:			Union[str, Any]											=	self._fs

        # use the -SSSSS-of-NNNNN pattern
        def _rename_shard(
            _A      :      int		,   _A      :      int		,   _A      :      int		,   ):
         rename(
             _A		,   fpath.replace('SSSSS'		,   F"""{shard_id:05d}"""	).replace('TTTTT'		,   F"""{task_id:05d}"""	)		,   fpath.replace('TTTTT-SSSSS'		,   F"""{global_shard_id:05d}"""	).replace('NNNNN'		,   F"""{total_shards:05d}"""	)		,   )

        snake_case_		:			Dict											=	[]
        snake_case_		:			List[str]											=	0
        for i in range(len(_A	)	):
         snake_case_		:			Union[str, Any]											=	task_id_and_num_shards[i]
         for shard_id in range(_A	):
          args.append([task_id, shard_id, global_shard_id]	)
          global_shard_id += 1
        self._spark.sparkContext.parallelize(_A		,   len(_A	)	).map(lambda _A	: _rename_shard(*_A	)	).collect()
       else:
        # don't use any pattern
        snake_case_		:			str											=	0
        snake_case_		:			Optional[Any]											=	task_id_and_num_shards[0][0]
        self._rename(
            fpath.replace('SSSSS'		,   F"""{shard_id:05d}"""	).replace('TTTTT'		,   F"""{task_id:05d}"""	)		,   fpath.replace(_A		,   ''	)		,   )






      def UpperCAmelCase_				(      self      :      Optional[int]		,   _A      :      "datasets.SplitGenerator"		,   )							->    SparkExamplesIterable:







       """simple docstring"""


       return SparkExamplesIterable(self.df	)








import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()
class SCREAMING_SNAKE_CASE_						(					snake_case_       ,				unittest.TestCase					):
						__magic_name__:			Dict       =       PriorTransformer
						__magic_name__:			str       =       "hidden_states"






						@property
						def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    Optional[int]:







							"""simple docstring"""


							snake_case_		:			Any											=	4
							snake_case_		:			int											=	8
							snake_case_		:			Dict											=	7

							snake_case_		:			Union[str, Any]											=	floats_tensor((batch_size, embedding_dim)	).to(_A	)

							snake_case_		:			int											=	floats_tensor((batch_size, embedding_dim)	).to(_A	)
							snake_case_		:			str											=	floats_tensor((batch_size, num_embeddings, embedding_dim)	).to(_A	)

							return {
							    "hidden_states": hidden_states,
							    "timestep": 2,
							    "proj_embedding": proj_embedding,
							    "encoder_hidden_states": encoder_hidden_states,
							}






						def UpperCAmelCase_				(      self      :      List[Any]		,   _A      :      List[Any]=0	)							->    str:







							"""simple docstring"""


							torch.manual_seed(_A	)
							snake_case_		:			List[Any]											=	4
							snake_case_		:			str											=	8
							snake_case_		:			Any											=	7

							snake_case_		:			List[Any]											=	torch.randn((batch_size, embedding_dim)	).to(_A	)

							snake_case_		:			Optional[int]											=	torch.randn((batch_size, embedding_dim)	).to(_A	)
							snake_case_		:			List[Any]											=	torch.randn((batch_size, num_embeddings, embedding_dim)	).to(_A	)

							return {
							    "hidden_states": hidden_states,
							    "timestep": 2,
							    "proj_embedding": proj_embedding,
							    "encoder_hidden_states": encoder_hidden_states,
							}






						@property
						def UpperCAmelCase_				(      self      :      str	)							->    Optional[Any]:







							"""simple docstring"""


							return (4, 8)






						@property
						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    List[Any]:







							"""simple docstring"""


							return (4, 8)






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







							"""simple docstring"""


							snake_case_		:			Union[str, Any]											=	{
							    'num_attention_heads': 2,
							    'attention_head_dim': 4,
							    'num_layers': 2,
							    'embedding_dim': 8,
							    'num_embeddings': 7,
							    'additional_embeddings': 4,
							}
							snake_case_		:			Tuple											=	self.dummy_input
							return init_dict, inputs_dict






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







							"""simple docstring"""


							snake_case_						,snake_case_		:			str											=	PriorTransformer.from_pretrained(
							    'hf-internal-testing/prior-dummy'		,   output_loading_info=_A	)
							self.assertIsNotNone(_A	)
							self.assertEqual(len(loading_info['missing_keys']	)		,   0	)

							model.to(_A	)
							snake_case_		:			Optional[Any]											=	model(**self.dummy_input	)[0]

							assert hidden_states is not None, "Make sure output is not None"






						def UpperCAmelCase_				(      self      :      Tuple	)							->    List[Any]:







							"""simple docstring"""


							snake_case_						,snake_case_		:			Optional[int]											=	self.prepare_init_args_and_inputs_for_common()

							snake_case_		:			Tuple											=	self.model_class(**_A	)
							snake_case_		:			Any											=	inspect.signature(model.forward	)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							snake_case_		:			List[str]											=	[*signature.parameters.keys()]

							snake_case_		:			int											=	['hidden_states', 'timestep']
							self.assertListEqual(arg_names[:2]		,   _A	)






						def UpperCAmelCase_				(      self      :      Union[str, Any]	)							->    List[Any]:







							"""simple docstring"""


							snake_case_		:			str											=	PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy'	)
							snake_case_		:			str											=	model.to(_A	)

							if hasattr(_A		,   'set_default_attn_processor'	):
								model.set_default_attn_processor()

							snake_case_		:			Optional[int]											=	self.get_dummy_seed_input()

							with torch.no_grad():
								snake_case_		:			Any											=	model(**_A	)[0]

							snake_case_		:			Any											=	output[0, :5].flatten().cpu()
							print(_A	)

							# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
							# the expected output slices are not the same for CPU and GPU.
							snake_case_		:			str											=	torch.tensor([-1.3_4_3_6, -0.2_8_7_0, 0.7_5_3_8, 0.4_3_6_8, -0.0_2_3_9]	)
							self.assertTrue(torch_all_close(_A		,   _A		,   rtol=1E-2	)	)
@slow
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






						def UpperCAmelCase_				(      self      :      Optional[Any]		,   _A      :      int=1		,   _A      :      int=768		,   _A      :      str=77		,   _A      :      List[str]=0	)							->    Tuple:







							"""simple docstring"""


							torch.manual_seed(_A	)
							snake_case_		:			Dict											=	batch_size
							snake_case_		:			Any											=	embedding_dim
							snake_case_		:			int											=	num_embeddings

							snake_case_		:			Dict											=	torch.randn((batch_size, embedding_dim)	).to(_A	)

							snake_case_		:			List[Any]											=	torch.randn((batch_size, embedding_dim)	).to(_A	)
							snake_case_		:			Optional[int]											=	torch.randn((batch_size, num_embeddings, embedding_dim)	).to(_A	)

							return {
							    "hidden_states": hidden_states,
							    "timestep": 2,
							    "proj_embedding": proj_embedding,
							    "encoder_hidden_states": encoder_hidden_states,
							}






						def UpperCAmelCase_				(      self      :      List[Any]	)							->    Optional[int]:







							"""simple docstring"""


							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






						@parameterized.expand(
						    [
						        # fmt: off
						        [13, [-0.5_8_6_1,  0.1_2_8_3, -0.0_9_3_1,  0.0_8_8_2,  0.4_4_7_6,  0.1_3_2_9, -0.0_4_9_8,  0.0_6_4_0]],
						        [37, [-0.4_9_1_3,  0.0_1_1_0, -0.0_4_8_3,  0.0_5_4_1,  0.4_9_5_4, -0.0_1_7_0,  0.0_3_5_4,  0.1_6_5_1]],
						        # fmt: on
						    ]	)
						def UpperCAmelCase_				(      self      :      Tuple		,   _A      :      List[Any]		,   _A      :      List[str]	)							->    Optional[int]:







							"""simple docstring"""


							snake_case_		:			str											=	PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'		,   subfolder='prior'	)
							model.to(_A	)
							snake_case_		:			Optional[Any]											=	self.get_dummy_seed_input(seed=_A	)

							with torch.no_grad():
								snake_case_		:			str											=	model(**_A	)[0]

							assert list(sample.shape	) == [1, 768]

							snake_case_		:			Optional[Any]											=	sample[0, :8].flatten().cpu()
							print(_A	)
							snake_case_		:			int											=	torch.tensor(_A	)

							assert torch_all_close(_A		,   _A		,   atol=1E-3	)

