



"""simple docstring"""






import collections
import inspect
import unittest
from typing import Dict, List, Tuple

from transformers import MaskFormerSwinConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, torch_device
from transformers.utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import MaskFormerSwinBackbone
   from transformers.models.maskformer import MaskFormerSwinModel






class 	__snake_case	:




  def __init__(			self					:						List[str]						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						Optional[int]=1_3						, __lowerCAmelCase					:						str=3_2						, __lowerCAmelCase					:						Optional[int]=2						, __lowerCAmelCase					:						Optional[int]=3						, __lowerCAmelCase					:						List[Any]=1_6						, __lowerCAmelCase					:						Dict=[1, 2, 1]						, __lowerCAmelCase					:						str=[2, 2, 4]						, __lowerCAmelCase					:						str=2						, __lowerCAmelCase					:						Optional[Any]=2.0						, __lowerCAmelCase					:						Union[str, Any]=True						, __lowerCAmelCase					:						Any=0.0						, __lowerCAmelCase					:						Union[str, Any]=0.0						, __lowerCAmelCase					:						str=0.1						, __lowerCAmelCase					:						List[Any]="gelu"						, __lowerCAmelCase					:						Any=False						, __lowerCAmelCase					:						str=True						, __lowerCAmelCase					:						List[Any]=0.02						, __lowerCAmelCase					:						Any=1E-5						, __lowerCAmelCase					:						List[str]=True						, __lowerCAmelCase					:						Union[str, Any]=None						, __lowerCAmelCase					:						Dict=True						, __lowerCAmelCase					:						Optional[int]=1_0						, __lowerCAmelCase					:						Union[str, Any]=8						, __lowerCAmelCase					:						List[Any]=["stage1", "stage2", "stage3"]						, __lowerCAmelCase					:						str=[1, 2, 3]						, ):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			parent
         _lowerCamelCase						:    Optional[int]         =			batch_size
         _lowerCamelCase						:    Tuple         =			image_size
         _lowerCamelCase						:    int         =			patch_size
         _lowerCamelCase						:    str         =			num_channels
         _lowerCamelCase						:    Tuple         =			embed_dim
         _lowerCamelCase						:    Optional[int]         =			depths
         _lowerCamelCase						:    List[Any]         =			num_heads
         _lowerCamelCase						:    str         =			window_size
         _lowerCamelCase						:    int         =			mlp_ratio
         _lowerCamelCase						:    Union[str, Any]         =			qkv_bias
         _lowerCamelCase						:    Dict         =			hidden_dropout_prob
         _lowerCamelCase						:    str         =			attention_probs_dropout_prob
         _lowerCamelCase						:    Union[str, Any]         =			drop_path_rate
         _lowerCamelCase						:    Union[str, Any]         =			hidden_act
         _lowerCamelCase						:    int         =			use_absolute_embeddings
         _lowerCamelCase						:    Union[str, Any]         =			patch_norm
         _lowerCamelCase						:    List[Any]         =			layer_norm_eps
         _lowerCamelCase						:    Optional[int]         =			initializer_range
         _lowerCamelCase						:    Tuple         =			is_training
         _lowerCamelCase						:    str         =			scope
         _lowerCamelCase						:    Optional[int]         =			use_labels
         _lowerCamelCase						:    Optional[Any]         =			type_sequence_label_size
         _lowerCamelCase						:    List[Any]         =			encoder_stride
         _lowerCamelCase						:    Dict         =			out_features
         _lowerCamelCase						:    Tuple         =			out_indices




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

         _lowerCamelCase						:    Union[str, Any]         =			None
         if self.use_labels:
                _lowerCamelCase						:    List[Any]         =			ids_tensor([self.batch_size]						, self.type_sequence_label_size		)

         _lowerCamelCase						:    Dict         =			self.get_config()

         return config, pixel_values, labels




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         return MaskFormerSwinConfig(
             image_size=self.image_size						, patch_size=self.patch_size						, num_channels=self.num_channels						, embed_dim=self.embed_dim						, depths=self.depths						, num_heads=self.num_heads						, window_size=self.window_size						, mlp_ratio=self.mlp_ratio						, qkv_bias=self.qkv_bias						, hidden_dropout_prob=self.hidden_dropout_prob						, attention_probs_dropout_prob=self.attention_probs_dropout_prob						, drop_path_rate=self.drop_path_rate						, hidden_act=self.hidden_act						, use_absolute_embeddings=self.use_absolute_embeddings						, path_norm=self.patch_norm						, layer_norm_eps=self.layer_norm_eps						, initializer_range=self.initializer_range						, encoder_stride=self.encoder_stride						, out_features=self.out_features						, out_indices=self.out_indices						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			MaskFormerSwinModel(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         _lowerCamelCase						:    Optional[Any]         =			model(__lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths		) - 1))
         _lowerCamelCase						:    Optional[Any]         =			int(config.embed_dim * 2 ** (len(config.depths		) - 1)		)

         self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, expected_seq_len, expected_dim)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						Dict		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			MaskFormerSwinBackbone(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         _lowerCamelCase						:    Optional[int]         =			model(__lowerCAmelCase		)

         # verify feature maps
         self.parent.assertEqual(len(result.feature_maps		)						, len(config.out_features		)		)
         self.parent.assertListEqual(list(result.feature_maps[0].shape		)						, [1_3, 1_6, 1_6, 1_6]		)

         # verify channels
         self.parent.assertEqual(len(model.channels		)						, len(config.out_features		)		)
         self.parent.assertListEqual(model.channels						, [1_6, 3_2, 6_4]		)

         # verify ValueError
         with self.parent.assertRaises(__lowerCAmelCase		):
                _lowerCamelCase						:    Union[str, Any]         =			['''stem''']
                _lowerCamelCase						:    List[Any]         =			MaskFormerSwinBackbone(config=__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			self.prepare_config_and_inputs()
         _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    Dict         =			config_and_inputs
         _lowerCamelCase						:    Union[str, Any]         =			{'''pixel_values''': pixel_values}
         return config, inputs_dict






@require_torch
class 	__snake_case	(     _lowercase  ,					_lowercase  ,					unittest.TestCase):
  snake_case__    :       Tuple        =						(
      (
          MaskFormerSwinModel,
          MaskFormerSwinBackbone,
      )
      if is_torch_available()
      else ()
  )
  snake_case__    :       Dict        =						{"feature-extraction": MaskFormerSwinModel} if is_torch_available() else {}
  snake_case__    :       str        =						False
  snake_case__    :       Union[str, Any]        =						False
  snake_case__    :       Optional[Any]        =						False
  snake_case__    :       int        =						False
  snake_case__    :       Union[str, Any]        =						False




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			MaskFormerSwinModelTester(self		)
         _lowerCamelCase						:    str         =			ConfigTester(self						, config_class=__lowerCAmelCase						, embed_dim=3_7		)




  @require_torch_multi_gpu
  @unittest.skip(
      reason=(
          '''`MaskFormerSwinModel` outputs `hidden_states_spatial_dimensions` which doesn\'t work well with'''
          ''' `nn.DataParallel`'''
      )		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         pass




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         self.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.create_and_test_config_with_num_labels()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         return




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_backbone(*__lowerCAmelCase		)




  @unittest.skip('''Swin does not use inputs_embeds'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         pass




  @unittest.skip('''Swin does not support feedforward chunking'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         pass




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Optional[Any]         =			self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
                _lowerCamelCase						:    Tuple         =			model_class(__lowerCAmelCase		)
                self.assertIsInstance(model.get_input_embeddings()						, (nn.Module)		)
                _lowerCamelCase						:    str         =			model.get_output_embeddings()
                self.assertTrue(x is None or isinstance(__lowerCAmelCase						, nn.Linear		)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
                _lowerCamelCase						:    Dict         =			model_class(__lowerCAmelCase		)
                _lowerCamelCase						:    Tuple         =			inspect.signature(model.forward		)
                # signature.parameters is an OrderedDict => so arg_names order is deterministic
                _lowerCamelCase						:    int         =			[*signature.parameters.keys()]

                _lowerCamelCase						:    Tuple         =			['''pixel_values''']
                self.assertListEqual(arg_names[:1]						, __lowerCAmelCase		)




  @unittest.skip(reason='''MaskFormerSwin is only used as backbone and doesn\'t support output_attentions'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         pass




  @unittest.skip(reason='''MaskFormerSwin is only used as an internal backbone'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         pass




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			model_class(__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()

         with torch.no_grad():
                _lowerCamelCase						:    List[Any]         =			model(**self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase		)		)

         _lowerCamelCase						:    Dict         =			outputs.hidden_states

         _lowerCamelCase						:    int         =			getattr(
             self.model_tester						, '''expected_num_hidden_layers'''						, len(self.model_tester.depths		) + 1		)
         self.assertEqual(len(__lowerCAmelCase		)						, __lowerCAmelCase		)

         # Swin has a different seq_length
         _lowerCamelCase						:    Dict         =			(
             config.patch_size
             if isinstance(config.patch_size						, collections.abc.Iterable		)
             else (config.patch_size, config.patch_size)
         )

         _lowerCamelCase						:    List[Any]         =			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

         self.assertListEqual(
             list(hidden_states[0].shape[-2:]		)						, [num_patches, self.model_tester.embed_dim]						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    List[Any]         =			self.model_tester.prepare_config_and_inputs_for_common()

         _lowerCamelCase						:    Tuple         =			(
             self.model_tester.image_size
             if isinstance(self.model_tester.image_size						, collections.abc.Iterable		)
             else (self.model_tester.image_size, self.model_tester.image_size)
         )

         for model_class in self.all_model_classes:
                _lowerCamelCase						:    Dict         =			True
                self.check_hidden_states_output(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                _lowerCamelCase						:    Any         =			True

                self.check_hidden_states_output(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    str         =			self.model_tester.prepare_config_and_inputs_for_common()
         _lowerCamelCase						:    List[str]         =			3

         _lowerCamelCase						:    Optional[Any]         =			(
             self.model_tester.image_size
             if isinstance(self.model_tester.image_size						, collections.abc.Iterable		)
             else (self.model_tester.image_size, self.model_tester.image_size)
         )
         _lowerCamelCase						:    Any         =			(
             config.patch_size
             if isinstance(config.patch_size						, collections.abc.Iterable		)
             else (config.patch_size, config.patch_size)
         )

         _lowerCamelCase						:    List[Any]         =			image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
         _lowerCamelCase						:    int         =			image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

         for model_class in self.all_model_classes:
                _lowerCamelCase						:    List[Any]         =			True
                self.check_hidden_states_output(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, (padded_height, padded_width)		)

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                _lowerCamelCase						:    Union[str, Any]         =			True
                self.check_hidden_states_output(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, (padded_height, padded_width)		)




  @unittest.skip(reason='''MaskFormerSwin doesn\'t have pretrained checkpoints'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         pass




  @unittest.skip(reason='''This will be fixed once MaskFormerSwin is replaced by native Swin'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         pass




  @unittest.skip(reason='''This will be fixed once MaskFormerSwin is replaced by native Swin'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         pass




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Optional[Any]         =			self.model_tester.prepare_config_and_inputs_for_common()

         def set_nan_tensor_to_zero(__lowerCAmelCase					:						Optional[int]		):
                _lowerCamelCase						:    str         =			0
                return t

         def check_equivalence(__lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						Optional[int]={}		):
                with torch.no_grad():
                       _lowerCamelCase						:    Any         =			model(**__lowerCAmelCase						, return_dict=__lowerCAmelCase						, **__lowerCAmelCase		)
                       _lowerCamelCase						:    Optional[Any]         =			model(**__lowerCAmelCase						, return_dict=__lowerCAmelCase						, **__lowerCAmelCase		).to_tuple()

                       def recursive_check(__lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						Any		):
                              if isinstance(__lowerCAmelCase						, (List, Tuple)		):
                                     for tuple_iterable_value, dict_iterable_value in zip(__lowerCAmelCase						, __lowerCAmelCase		):
                                            recursive_check(__lowerCAmelCase						, __lowerCAmelCase		)
                              elif isinstance(__lowerCAmelCase						, __lowerCAmelCase		):
                                     for tuple_iterable_value, dict_iterable_value in zip(
                                         tuple_object.values()						, dict_object.values()		):
                                            recursive_check(__lowerCAmelCase						, __lowerCAmelCase		)
                              elif tuple_object is None:
                                     return
                              else:
                                     self.assertTrue(
                                         torch.allclose(
                                             set_nan_tensor_to_zero(__lowerCAmelCase		)						, set_nan_tensor_to_zero(__lowerCAmelCase		)						, atol=1E-5		)						, msg=(
                                             '''Tuple and dict output are not equal. Difference:'''
                                             f''' {torch.max(torch.abs(tuple_object - dict_object		)		)}. Tuple has `nan`:'''
                                             f''' {torch.isnan(__lowerCAmelCase		).any()} and `inf`: {torch.isinf(__lowerCAmelCase		)}. Dict has'''
                                             f''' `nan`: {torch.isnan(__lowerCAmelCase		).any()} and `inf`: {torch.isinf(__lowerCAmelCase		)}.'''
                                         )						, )

                       recursive_check(__lowerCAmelCase						, __lowerCAmelCase		)

         for model_class in self.all_model_classes:
                _lowerCamelCase						:    Any         =			model_class(__lowerCAmelCase		)
                model.to(__lowerCAmelCase		)
                model.eval()

                _lowerCamelCase						:    int         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase		)
                _lowerCamelCase						:    List[str]         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase		)
                check_equivalence(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)

                _lowerCamelCase						:    int         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)
                _lowerCamelCase						:    str         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)
                check_equivalence(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)

                _lowerCamelCase						:    Optional[int]         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase		)
                _lowerCamelCase						:    int         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase		)
                check_equivalence(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, {'''output_hidden_states''': True}		)

                _lowerCamelCase						:    List[str]         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)
                _lowerCamelCase						:    Optional[int]         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)
                check_equivalence(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, {'''output_hidden_states''': True}		)






@require_torch
class 	__snake_case	(     unittest.TestCase  ,					_lowercase):
  snake_case__    :       int        =						(MaskFormerSwinBackbone,) if is_torch_available() else ()
  snake_case__    :       List[str]        =						MaskFormerSwinConfig




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			MaskFormerSwinModelTester(self		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Optional[Any]         =			self.model_tester.prepare_config_and_inputs_for_common()
         _lowerCamelCase						:    Union[str, Any]         =			inputs_dict['''pixel_values'''].shape[0]

         for backbone_class in self.all_model_classes:
                _lowerCamelCase						:    Optional[int]         =			backbone_class(__lowerCAmelCase		)
                backbone.to(__lowerCAmelCase		)
                backbone.eval()

                _lowerCamelCase						:    int         =			backbone(**__lowerCAmelCase		)

                # Test default outputs and verify feature maps
                self.assertIsInstance(outputs.feature_maps						, __lowerCAmelCase		)
                self.assertTrue(len(outputs.feature_maps		) == len(backbone.channels		)		)
                for feature_map, n_channels in zip(outputs.feature_maps						, backbone.channels		):
                       self.assertTrue(feature_map.shape[:2]						, (batch_size, n_channels)		)
                self.assertIsNone(outputs.hidden_states		)
                self.assertIsNone(outputs.attentions		)

                # Test output_hidden_states=True
                _lowerCamelCase						:    str         =			backbone(**__lowerCAmelCase						, output_hidden_states=__lowerCAmelCase		)
                self.assertIsNotNone(outputs.hidden_states		)
                self.assertTrue(len(outputs.hidden_states		)						, len(backbone.stage_names		)		)
                # We skip the stem layer
                for hidden_states, n_channels in zip(outputs.hidden_states[1:]						, backbone.channels		):
                       for hidden_state in hidden_states:
                              # Hidden states are in the format (batch_size, (height * width), n_channels)
                              _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    List[str]         =			hidden_state.shape
                              self.assertTrue((h_batch_size, h_n_channels)						, (batch_size, n_channels)		)

            # Test output_attentions=True
                if self.has_attentions:
                       _lowerCamelCase						:    List[str]         =			backbone(**__lowerCAmelCase						, output_attentions=__lowerCAmelCase		)
                       self.assertIsNotNone(outputs.attentions		)







"""simple docstring"""






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


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

if TYPE_CHECKING:
   import pyspark






@dataclass
class 	__snake_case	(     datasets.BuilderConfig):
  snake_case__    :       Optional[datasets.Features]        =						None






def 	snake_case_  (				A_ :		"pyspark.sql.DataFrame",    A_ :		List[int],    ):


       '''simple docstring'''





       import pyspark

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

       return generate_fn






class 	__snake_case	(     _BaseExamplesIterable):




  def __init__(			self					:						Tuple						, __lowerCAmelCase					:						"pyspark.sql.DataFrame"						, __lowerCAmelCase					:						Optional[int]=None						, ):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			df
         _lowerCamelCase						:    Union[str, Any]         =			partition_order or range(self.df.rdd.getNumPartitions()		)
         _lowerCamelCase						:    Dict         =			_generate_iterable_examples(self.df						, self.partition_order		)




  def __iter__(			self					:						List[Any]		):






         """simple docstring"""






         yield from self.generate_examples_fn()




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						np.random.Generator		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			list(range(self.df.rdd.getNumPartitions()		)		)
         generator.shuffle(__lowerCAmelCase		)
         return SparkExamplesIterable(self.df						, partition_order=__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			self.split_shard_indices_by_worker(__lowerCAmelCase						, __lowerCAmelCase		)
         return SparkExamplesIterable(self.df						, partition_order=__lowerCAmelCase		)




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         return len(self.partition_order		)






class 	__snake_case	(     datasets.DatasetBuilder):
  snake_case__    :       List[Any]        =						SparkConfig




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






         """simple docstring"""






         import pyspark

         _lowerCamelCase						:    Optional[int]         =			pyspark.sql.SparkSession.builder.getOrCreate()
         _lowerCamelCase						:    int         =			df
         _lowerCamelCase						:    Any         =			working_dir

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         def create_cache_and_write_probe(__lowerCAmelCase					:						Optional[int]		):
                # 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=__lowerCAmelCase		)
                _lowerCamelCase						:    Optional[int]         =			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(__lowerCAmelCase						, '''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:
                _lowerCamelCase						:    Optional[Any]         =			(
                    self._spark.sparkContext.parallelize(range(1		)						, 1		).mapPartitions(__lowerCAmelCase		).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 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






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




  def 						SCREAMING_SNAKE_CASE    (			self					:						int						, __lowerCAmelCase					:						datasets.download.download_manager.DownloadManager		):






         """simple docstring"""






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




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						Union[str, Any]		):






         """simple docstring"""






         import pyspark

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

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						int						, ):






         """simple docstring"""






         import pyspark

         _lowerCamelCase						:    Optional[Any]         =			ParquetWriter if file_format == '''parquet''' else ArrowWriter
         _lowerCamelCase						:    List[Any]         =			os.path.join(self._working_dir						, os.path.basename(__lowerCAmelCase		)		) if self._working_dir else fpath
         _lowerCamelCase						:    Dict         =			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.
         _lowerCamelCase						:    str         =			self.config.features
         _lowerCamelCase						:    Dict         =			self._writer_batch_size
         _lowerCamelCase						:    List[str]         =			self._fs.storage_options

         def write_arrow(__lowerCAmelCase					:						List[str]		):
                # Within the same SparkContext, no two task attempts will share the same attempt ID.
                _lowerCamelCase						:    List[str]         =			pyspark.TaskContext().taskAttemptId()
                _lowerCamelCase						:    Any         =			next(__lowerCAmelCase						, __lowerCAmelCase		)
                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''']						, )
                _lowerCamelCase						:    List[Any]         =			0
                _lowerCamelCase						:    Optional[int]         =			writer_class(
                    features=__lowerCAmelCase						, path=working_fpath.replace('''SSSSS'''						, f'''{shard_id:05d}'''		).replace('''TTTTT'''						, f'''{task_id:05d}'''		)						, writer_batch_size=__lowerCAmelCase						, storage_options=__lowerCAmelCase						, embed_local_files=__lowerCAmelCase						, )
                _lowerCamelCase						:    int         =			pa.Table.from_batches([first_batch]		)
                writer.write_table(__lowerCAmelCase		)
                for batch in it:
                       if max_shard_size is not None and writer._num_bytes >= max_shard_size:
                              _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			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
                              _lowerCamelCase						:    Optional[int]         =			writer_class(
                                  features=writer._features						, path=working_fpath.replace('''SSSSS'''						, f'''{shard_id:05d}'''		).replace('''TTTTT'''						, f'''{task_id:05d}'''		)						, writer_batch_size=__lowerCAmelCase						, storage_options=__lowerCAmelCase						, embed_local_files=__lowerCAmelCase						, )
                       _lowerCamelCase						:    Optional[int]         =			pa.Table.from_batches([batch]		)
                       writer.write_table(__lowerCAmelCase		)

                if writer._num_bytes > 0:
                       _lowerCamelCase      ,   _lowerCamelCase						:    Optional[int]         =			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(__lowerCAmelCase		)		):
                              _lowerCamelCase						:    Optional[Any]         =			os.path.join(os.path.dirname(__lowerCAmelCase		)						, os.path.basename(__lowerCAmelCase		)		)
                              shutil.move(__lowerCAmelCase						, __lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			(
             self.df.mapInArrow(__lowerCAmelCase						, '''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 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						"datasets.SplitGenerator"						, __lowerCAmelCase					:						str = "arrow"						, __lowerCAmelCase					:						Optional[Union[str, int]] = None						, __lowerCAmelCase					:						Optional[int] = None						, **__lowerCAmelCase					:						Tuple						, ):






         """simple docstring"""






         self._validate_cache_dir()

         _lowerCamelCase						:    str         =			convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE		)
         self._repartition_df_if_needed(__lowerCAmelCase		)
         _lowerCamelCase						:    str         =			not is_remote_filesystem(self._fs		)
         _lowerCamelCase						:    Tuple         =			os.path.join if is_local else posixpath.join

         _lowerCamelCase						:    int         =			'''-TTTTT-SSSSS-of-NNNNN'''
         _lowerCamelCase						:    Tuple         =			f'''{self.name}-{split_generator.name}{SUFFIX}.{file_format}'''
         _lowerCamelCase						:    List[Any]         =			path_join(self._output_dir						, __lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			0
         _lowerCamelCase						:    Any         =			0
         _lowerCamelCase						:    str         =			0
         _lowerCamelCase						:    int         =			[]
         _lowerCamelCase						:    List[str]         =			[]

         for task_id, content in self._prepare_split_single(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		):
                (
                    (
                    _lowerCamelCase
                )      ,   (
                    _lowerCamelCase
                )      ,   (
                    _lowerCamelCase
                )      ,   (
                    _lowerCamelCase
                )      ,   
                )						:    str         =			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(__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			total_num_examples
         _lowerCamelCase						:    str         =			total_num_bytes

         # should rename everything at the end
         logger.debug(f'''Renaming {total_shards} shards.'''		)
         if total_shards > 1:
                _lowerCamelCase						:    Optional[Any]         =			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.
                _lowerCamelCase						:    str         =			self._fs

                # use the -SSSSS-of-NNNNN pattern
                def _rename_shard(
                    __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, ):
                       rename(
                           __lowerCAmelCase						, 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}'''		)						, )

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						"datasets.SplitGenerator"						, ):






         """simple docstring"""






         return SparkExamplesIterable(self.df		)


