import os
import tempfile
import unittest

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       FlaubertForMultipleChoice,
       FlaubertForQuestionAnswering,
       FlaubertForQuestionAnsweringSimple,
       FlaubertForSequenceClassification,
       FlaubertForTokenClassification,
       FlaubertModel,
       FlaubertWithLMHeadModel,
   )
   from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST





class 		_lowercase						(					snake_case_       ):

 def __init__( self	:							int   ,      snake_case	:							Optional[int]   ,      snake_case	:							Tuple=1_3   ,      snake_case	:							Tuple=7   ,      snake_case	:							Any=True   ,      snake_case	:							Tuple=True   ,      snake_case	:							Tuple=True   ,      snake_case	:							Any=True   ,      snake_case	:							Optional[int]=True   ,      snake_case	:							int=False   ,      snake_case	:							Any=False   ,      snake_case	:							Dict=False   ,      snake_case	:							Optional[int]=2   ,      snake_case	:							str=9_9   ,      snake_case	:							Tuple=0   ,      snake_case	:							Any=3_2   ,      snake_case	:							List[Any]=5   ,      snake_case	:							Any=4   ,      snake_case	:							Optional[int]=0.1   ,      snake_case	:							str=0.1   ,      snake_case	:							int=5_1_2   ,      snake_case	:							int=1_2   ,      snake_case	:							Tuple=2   ,      snake_case	:							Tuple=0.02   ,      snake_case	:							Any=3   ,      snake_case	:							Dict=4   ,      snake_case	:							int="last"   ,      snake_case	:							Dict=None   ,      snake_case	:							Any=None   ,      )   ->						Tuple:

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					parent
  UpperCamelCase_					:						Union[str, Any] =					batch_size
  UpperCamelCase_					:						Optional[Any] =					seq_length
  UpperCamelCase_					:						List[str] =					is_training
  UpperCamelCase_					:						str =					use_input_lengths
  UpperCamelCase_					:						Optional[int] =					use_token_type_ids
  UpperCamelCase_					:						Optional[Any] =					use_labels
  UpperCamelCase_					:						int =					gelu_activation
  UpperCamelCase_					:						str =					sinusoidal_embeddings
  UpperCamelCase_					:						Any =					causal
  UpperCamelCase_					:						Dict =					asm
  UpperCamelCase_					:						Union[str, Any] =					n_langs
  UpperCamelCase_					:						Optional[int] =					vocab_size
  UpperCamelCase_					:						Dict =					n_special
  UpperCamelCase_					:						str =					hidden_size
  UpperCamelCase_					:						List[str] =					num_hidden_layers
  UpperCamelCase_					:						List[str] =					num_attention_heads
  UpperCamelCase_					:						int =					hidden_dropout_prob
  UpperCamelCase_					:						int =					attention_probs_dropout_prob
  UpperCamelCase_					:						Any =					max_position_embeddings
  UpperCamelCase_					:						Dict =					type_vocab_size
  UpperCamelCase_					:						List[Any] =					type_sequence_label_size
  UpperCamelCase_					:						int =					initializer_range
  UpperCamelCase_					:						Optional[int] =					num_labels
  UpperCamelCase_					:						Optional[int] =					num_choices
  UpperCamelCase_					:						List[Any] =					summary_type
  UpperCamelCase_					:						List[str] =					use_proj
  UpperCamelCase_					:						str =					scope

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

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					ids_tensor([self.batch_size, self.seq_length]   ,      self.vocab_size			)
  UpperCamelCase_					:						str =					random_attention_mask([self.batch_size, self.seq_length]			)

  UpperCamelCase_					:						List[Any] =					None
  if self.use_input_lengths:
   UpperCamelCase_					:						Optional[int] =					(
       ids_tensor([self.batch_size]   ,      vocab_size=2			) + self.seq_length - 2
   )  # small variation of seq_length

  UpperCamelCase_					:						Dict =					None
  if self.use_token_type_ids:
   UpperCamelCase_					:						Dict =					ids_tensor([self.batch_size, self.seq_length]   ,      self.n_langs			)

  UpperCamelCase_					:						Any =					None
  UpperCamelCase_					:						Optional[Any] =					None
  UpperCamelCase_					:						Optional[int] =					None
  if self.use_labels:
   UpperCamelCase_					:						Any =					ids_tensor([self.batch_size]   ,      self.type_sequence_label_size			)
   UpperCamelCase_					:						Union[str, Any] =					ids_tensor([self.batch_size, self.seq_length]   ,      self.num_labels			)
   UpperCamelCase_					:						List[str] =					ids_tensor([self.batch_size]   ,      2			).float()
   UpperCamelCase_					:						Tuple =					ids_tensor([self.batch_size]   ,      self.num_choices			)

  UpperCamelCase_					:						Optional[int] =					self.get_config()

  return (
      config,
      input_ids,
      token_type_ids,
      input_lengths,
      sequence_labels,
      token_labels,
      is_impossible_labels,
      choice_labels,
      input_mask,
  )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[str]			)   ->						int:

  """simple docstring"""





  return FlaubertConfig(
      vocab_size=self.vocab_size   ,      n_special=self.n_special   ,      emb_dim=self.hidden_size   ,      n_layers=self.num_hidden_layers   ,      n_heads=self.num_attention_heads   ,      dropout=self.hidden_dropout_prob   ,      attention_dropout=self.attention_probs_dropout_prob   ,      gelu_activation=self.gelu_activation   ,      sinusoidal_embeddings=self.sinusoidal_embeddings   ,      asm=self.asm   ,      causal=self.causal   ,      n_langs=self.n_langs   ,      max_position_embeddings=self.max_position_embeddings   ,      initializer_range=self.initializer_range   ,      summary_type=self.summary_type   ,      use_proj=self.use_proj   ,      )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any   ,      snake_case	:							Any   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							List[str]   ,      snake_case	:							Optional[Any]   ,      snake_case	:							str   ,      snake_case	:							Dict   ,      snake_case	:							List[str]   ,      snake_case	:							List[Any]   ,      snake_case	:							Tuple   ,      )   ->						Union[str, Any]:

  """simple docstring"""





  UpperCamelCase_					:						str =					FlaubertModel(config=snake_case			)
  model.to(snake_case			)
  model.eval()
  UpperCamelCase_					:						Any =					model(snake_case   ,      lengths=snake_case   ,      langs=snake_case			)
  UpperCamelCase_					:						List[Any] =					model(snake_case   ,      langs=snake_case			)
  UpperCamelCase_					:						Union[str, Any] =					model(snake_case			)
  self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size)			)

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

  """simple docstring"""





  UpperCamelCase_					:						Any =					FlaubertWithLMHeadModel(snake_case			)
  model.to(snake_case			)
  model.eval()

  UpperCamelCase_					:						str =					model(snake_case   ,      token_type_ids=snake_case   ,      labels=snake_case			)
  self.parent.assertEqual(result.loss.shape   ,      ()			)
  self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.vocab_size)			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[Any]   ,      snake_case	:							Optional[Any]   ,      snake_case	:							Dict   ,      snake_case	:							List[Any]   ,      snake_case	:							Tuple   ,      snake_case	:							Dict   ,      snake_case	:							Optional[Any]   ,      snake_case	:							str   ,      snake_case	:							Any   ,      snake_case	:							Tuple   ,      )   ->						Union[str, Any]:

  """simple docstring"""





  UpperCamelCase_					:						List[Any] =					FlaubertForQuestionAnsweringSimple(snake_case			)
  model.to(snake_case			)
  model.eval()

  UpperCamelCase_					:						List[str] =					model(snake_case			)

  UpperCamelCase_					:						Union[str, Any] =					model(snake_case   ,      start_positions=snake_case   ,      end_positions=snake_case			)
  self.parent.assertEqual(result.start_logits.shape   ,      (self.batch_size, self.seq_length)			)
  self.parent.assertEqual(result.end_logits.shape   ,      (self.batch_size, self.seq_length)			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							str   ,      snake_case	:							str   ,      snake_case	:							List[str]   ,      snake_case	:							Optional[Any]   ,      snake_case	:							Tuple   ,      snake_case	:							str   ,      snake_case	:							List[Any]   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							Dict   ,      snake_case	:							List[Any]   ,      )   ->						Optional[Any]:

  """simple docstring"""





  UpperCamelCase_					:						str =					FlaubertForQuestionAnswering(snake_case			)
  model.to(snake_case			)
  model.eval()

  UpperCamelCase_					:						str =					model(snake_case			)

  UpperCamelCase_					:						str =					model(
      snake_case   ,      start_positions=snake_case   ,      end_positions=snake_case   ,      cls_index=snake_case   ,      is_impossible=snake_case   ,      p_mask=snake_case   ,      )

  UpperCamelCase_					:						Optional[int] =					model(
      snake_case   ,      start_positions=snake_case   ,      end_positions=snake_case   ,      cls_index=snake_case   ,      is_impossible=snake_case   ,      )

  ((UpperCamelCase_), )					:						Optional[int] =					result_with_labels.to_tuple()

  UpperCamelCase_					:						Union[str, Any] =					model(snake_case   ,      start_positions=snake_case   ,      end_positions=snake_case			)

  ((UpperCamelCase_), )					:						List[str] =					result_with_labels.to_tuple()

  self.parent.assertEqual(result_with_labels.loss.shape   ,      ()			)
  self.parent.assertEqual(result.start_top_log_probs.shape   ,      (self.batch_size, model.config.start_n_top)			)
  self.parent.assertEqual(result.start_top_index.shape   ,      (self.batch_size, model.config.start_n_top)			)
  self.parent.assertEqual(
      result.end_top_log_probs.shape   ,      (self.batch_size, model.config.start_n_top * model.config.end_n_top)			)
  self.parent.assertEqual(
      result.end_top_index.shape   ,      (self.batch_size, model.config.start_n_top * model.config.end_n_top)			)
  self.parent.assertEqual(result.cls_logits.shape   ,      (self.batch_size,)			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[int]   ,      snake_case	:							Optional[int]   ,      snake_case	:							int   ,      snake_case	:							Any   ,      snake_case	:							Dict   ,      snake_case	:							int   ,      snake_case	:							str   ,      snake_case	:							Optional[int]   ,      snake_case	:							str   ,      snake_case	:							Dict   ,      )   ->						str:

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					FlaubertForSequenceClassification(snake_case			)
  model.to(snake_case			)
  model.eval()

  UpperCamelCase_					:						List[str] =					model(snake_case			)
  UpperCamelCase_					:						Any =					model(snake_case   ,      labels=snake_case			)

  self.parent.assertEqual(result.loss.shape   ,      ()			)
  self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.type_sequence_label_size)			)

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

  """simple docstring"""





  UpperCamelCase_					:						Dict =					self.num_labels
  UpperCamelCase_					:						int =					FlaubertForTokenClassification(snake_case			)
  model.to(snake_case			)
  model.eval()

  UpperCamelCase_					:						Optional[int] =					model(snake_case   ,      attention_mask=snake_case   ,      labels=snake_case			)
  self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.num_labels)			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]   ,      snake_case	:							Tuple   ,      snake_case	:							Optional[int]   ,      snake_case	:							Dict   ,      snake_case	:							Any   ,      snake_case	:							Tuple   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							List[Any]   ,      snake_case	:							Optional[int]   ,      snake_case	:							str   ,      )   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						str =					self.num_choices
  UpperCamelCase_					:						List[Any] =					FlaubertForMultipleChoice(config=snake_case			)
  model.to(snake_case			)
  model.eval()
  UpperCamelCase_					:						Dict =					input_ids.unsqueeze(1			).expand(-1   ,      self.num_choices   ,      -1			).contiguous()
  UpperCamelCase_					:						Optional[int] =					token_type_ids.unsqueeze(1			).expand(-1   ,      self.num_choices   ,      -1			).contiguous()
  UpperCamelCase_					:						int =					input_mask.unsqueeze(1			).expand(-1   ,      self.num_choices   ,      -1			).contiguous()
  UpperCamelCase_					:						Tuple =					model(
      snake_case   ,      attention_mask=snake_case   ,      token_type_ids=snake_case   ,      labels=snake_case   ,      )
  self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_choices)			)


 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[int]			)   ->						Optional[int]:

  """simple docstring"""





  UpperCamelCase_					:						Any =					self.prepare_config_and_inputs()
  (
      (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), 
  )					:						List[str] =					config_and_inputs
  UpperCamelCase_					:						Union[str, Any] =					{
      'input_ids': input_ids,
      'token_type_ids': token_type_ids,
      'lengths': input_lengths,
      'attention_mask': input_mask,
  }
  return config, inputs_dict





@require_torch
class 		_lowercase						(					snake_case_						,      snake_case_						,      unittest.TestCase       ):
 lowercase				     =   (
     (
         FlaubertModel,
         FlaubertWithLMHeadModel,
         FlaubertForQuestionAnswering,
         FlaubertForQuestionAnsweringSimple,
         FlaubertForSequenceClassification,
         FlaubertForTokenClassification,
         FlaubertForMultipleChoice,
     )
     if is_torch_available()
     else ()
 )
 lowercase				     =   (
     {
         'feature-extraction': FlaubertModel,
         'fill-mask': FlaubertWithLMHeadModel,
         'question-answering': FlaubertForQuestionAnsweringSimple,
         'text-classification': FlaubertForSequenceClassification,
         'token-classification': FlaubertForTokenClassification,
         'zero-shot': FlaubertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[int]   ,      snake_case	:							Dict   ,      snake_case	:							str   ,      snake_case	:							Any   ,      snake_case	:							Any   ,      snake_case	:							List[str]			)   ->						List[str]:

  """simple docstring"""





  if (
      pipeline_test_casse_name == "QAPipelineTests"
      and tokenizer_name is not None
      and not tokenizer_name.endswith('Fast'			)
  ):
   # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
   # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
   # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
   return True

  return False

 def 						SCREAMING_SNAKE_CASE__   ( self	:							str   ,      snake_case	:							List[Any]   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							Optional[Any]=False			)   ->						Optional[Any]:

  """simple docstring"""





  UpperCamelCase_					:						List[Any] =					super()._prepare_for_class(snake_case   ,      snake_case   ,      return_labels=snake_case			)

  if return_labels:
   if model_class.__name__ == "FlaubertForQuestionAnswering":
    UpperCamelCase_					:						Tuple =					torch.zeros(
        self.model_tester.batch_size   ,      dtype=torch.long   ,      device=snake_case			)
    UpperCamelCase_					:						Tuple =					torch.zeros(
        self.model_tester.batch_size   ,      dtype=torch.long   ,      device=snake_case			)

  return inputs_dict

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

  """simple docstring"""





  UpperCamelCase_					:						Dict =					FlaubertModelTester(self			)
  UpperCamelCase_					:						Dict =					ConfigTester(self   ,      config_class=snake_case   ,      emb_dim=3_7			)

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

  """simple docstring"""





  self.config_tester.run_common_tests()

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any			)   ->						str:

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_flaubert_model(*snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]			)   ->						Tuple:

  """simple docstring"""





  UpperCamelCase_					:						Dict =					self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_flaubert_lm_head(*snake_case			)

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

  """simple docstring"""





  UpperCamelCase_					:						List[Any] =					self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_flaubert_simple_qa(*snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Dict			)   ->						Dict:

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_flaubert_qa(*snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[int]			)   ->						Optional[int]:

  """simple docstring"""





  UpperCamelCase_					:						Dict =					self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_flaubert_sequence_classif(*snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						Optional[int]:

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_flaubert_token_classif(*snake_case			)

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

  """simple docstring"""





  UpperCamelCase_					:						int =					self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_flaubert_multiple_choice(*snake_case			)

 @slow
 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]			)   ->						Any:

  """simple docstring"""





  for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   UpperCamelCase_					:						List[str] =					FlaubertModel.from_pretrained(snake_case			)
   self.assertIsNotNone(snake_case			)


 @slow
 @require_torch_gpu
 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[str]			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_, UpperCamelCase_					:						Dict =					self.model_tester.prepare_config_and_inputs_for_common()
  for model_class in self.all_model_classes:
   # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
   if model_class == FlaubertForMultipleChoice:
    return

   UpperCamelCase_					:						List[str] =					True
   UpperCamelCase_					:						Any =					model_class(config=snake_case			)

   UpperCamelCase_					:						str =					self._prepare_for_class(snake_case   ,      snake_case			)
   UpperCamelCase_					:						List[str] =					torch.jit.trace(
       snake_case   ,      (inputs_dict['input_ids'].to('cpu'			), inputs_dict['attention_mask'].to('cpu'			))			)

   with tempfile.TemporaryDirectory() as tmp:
    torch.jit.save(snake_case   ,      os.path.join(snake_case   ,      'traced_model.pt'			)			)
    UpperCamelCase_					:						Union[str, Any] =					torch.jit.load(os.path.join(snake_case   ,      'traced_model.pt'			)   ,      map_location=snake_case			)
    loaded(inputs_dict['input_ids'].to(snake_case			)   ,      inputs_dict['attention_mask'].to(snake_case			)			)





@require_torch
class 		_lowercase						(					unittest.TestCase       ):

 @slow
 def 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						Union[str, Any] =					FlaubertModel.from_pretrained('flaubert/flaubert_base_cased'			)
  UpperCamelCase_					:						int =					torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]			)
  with torch.no_grad():
   UpperCamelCase_					:						Optional[Any] =					model(snake_case			)[0]
  UpperCamelCase_					:						Dict =					torch.Size((1, 1_1, 7_6_8)			)
  self.assertEqual(output.shape   ,      snake_case			)
  UpperCamelCase_					:						Any =					torch.tensor(
      [[[-2.6251, -1.4298, -0.0227], [-2.8510, -1.6387, 0.2258], [-2.8114, -1.1832, -0.3066]]]			)

  self.assertTrue(torch.allclose(output[:, :3, :3]   ,      snake_case   ,      atol=1e-4			)			)




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


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

if TYPE_CHECKING:
   import pyspark





@dataclass
class 		_lowercase						(					datasets.BuilderConfig       ):
 lowercase				     =   None





def    __lowercase     (						lowerCamelCase					:     "pyspark.sql.DataFrame"  ,  lowerCamelCase					:     List[int]  ,  ):
 import pyspark

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

 return generate_fn





class 		_lowercase						(					_BaseExamplesIterable       ):

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

  """simple docstring"""





  UpperCamelCase_					:						Dict =					df
  UpperCamelCase_					:						int =					partition_order or range(self.df.rdd.getNumPartitions()			)
  UpperCamelCase_					:						Optional[Any] =					_generate_iterable_examples(self.df   ,      self.partition_order			)

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

  """simple docstring"""





  yield from self.generate_examples_fn()

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any   ,      snake_case	:							np.random.Generator			)   ->						"SparkExamplesIterable":

  """simple docstring"""





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

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int   ,      snake_case	:							int   ,      snake_case	:							int			)   ->						"SparkExamplesIterable":

  """simple docstring"""





  UpperCamelCase_					:						Tuple =					self.split_shard_indices_by_worker(snake_case   ,      snake_case			)
  return SparkExamplesIterable(self.df   ,      partition_order=snake_case			)


 @property
 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						int:

  """simple docstring"""





  return len(self.partition_order			)





class 		_lowercase						(					datasets.DatasetBuilder       ):
 lowercase				     =   SparkConfig

 def __init__( self	:							List[Any]   ,      snake_case	:							"pyspark.sql.DataFrame"   ,      snake_case	:							str = None   ,      snake_case	:							str = None   ,      **snake_case	:							Optional[Any]   ,      )   ->						List[str]:

  """simple docstring"""





  import pyspark

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

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

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any			)   ->						Dict:

  """simple docstring"""





  def create_cache_and_write_probe(snake_case	:							str			):
   # 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=snake_case			)
   UpperCamelCase_					:						Tuple =					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(snake_case   ,      '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:
   UpperCamelCase_					:						Tuple =					(
       self._spark.sparkContext.parallelize(range(1			)   ,      1			).mapPartitions(snake_case			).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	:							int			)   ->						Union[str, Any]:

  """simple docstring"""





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

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]   ,      snake_case	:							datasets.download.download_manager.DownloadManager			)   ->						Optional[int]:

  """simple docstring"""





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

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]   ,      snake_case	:							Optional[int]			)   ->						List[Any]:

  """simple docstring"""





  import pyspark

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

  UpperCamelCase_					:						List[str] =					self.df.count()
  UpperCamelCase_					:						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.
  UpperCamelCase_					:						str =					(
      self.df.limit(snake_case			)
      .repartition(1			)
      .mapInArrow(snake_case   ,      'batch_bytes: long'			)
      .agg(pyspark.sql.functions.sum('batch_bytes'			).alias('sample_bytes'			)			)
      .collect()[0]
      .sample_bytes
      / sample_num_rows
  )
  UpperCamelCase_					:						Optional[int] =					approx_bytes_per_row * df_num_rows
  if approx_total_size > max_shard_size:
   # Make sure there is at least one row per partition.
   UpperCamelCase_					:						Optional[Any] =					min(snake_case   ,      int(approx_total_size / max_shard_size			)			)
   UpperCamelCase_					:						int =					self.df.repartition(snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]   ,      snake_case	:							str   ,      snake_case	:							str   ,      snake_case	:							int   ,      )   ->						Iterable[Tuple[int, bool, Union[int, tuple]]]:

  """simple docstring"""





  import pyspark

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

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

   if writer._num_bytes > 0:
    UpperCamelCase_, UpperCamelCase_					:						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(snake_case			)			):
     UpperCamelCase_					:						Dict =					os.path.join(os.path.dirname(snake_case			)   ,      os.path.basename(snake_case			)			)
     shutil.move(snake_case   ,      snake_case			)

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

  """simple docstring"""





  self._validate_cache_dir()

  UpperCamelCase_					:						Optional[int] =					convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE			)
  self._repartition_df_if_needed(snake_case			)
  UpperCamelCase_					:						List[str] =					not is_remote_filesystem(self._fs			)
  UpperCamelCase_					:						List[Any] =					os.path.join if is_local else posixpath.join

  UpperCamelCase_					:						Optional[int] =					'-TTTTT-SSSSS-of-NNNNN'
  UpperCamelCase_					:						Dict =					f"{self.name}-{split_generator.name}{SUFFIX}.{file_format}"
  UpperCamelCase_					:						int =					path_join(self._output_dir   ,      snake_case			)

  UpperCamelCase_					:						int =					0
  UpperCamelCase_					:						Optional[int] =					0
  UpperCamelCase_					:						Union[str, Any] =					0
  UpperCamelCase_					:						Optional[Any] =					[]
  UpperCamelCase_					:						Any =					[]

  for task_id, content in self._prepare_split_single(snake_case   ,      snake_case   ,      snake_case			):
   (
       (
       UpperCamelCase_
   ), (
       UpperCamelCase_
   ), (
       UpperCamelCase_
   ), (
       UpperCamelCase_
   ), 
   )					:						Optional[Any] =					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(snake_case			)

  UpperCamelCase_					:						Optional[Any] =					total_num_examples
  UpperCamelCase_					:						Any =					total_num_bytes

  # should rename everything at the end
  logger.debug(f"Renaming {total_shards} shards."			)
  if total_shards > 1:
   UpperCamelCase_					:						List[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.
   UpperCamelCase_					:						int =					self._fs

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

   UpperCamelCase_					:						Any =					[]
   UpperCamelCase_					:						Optional[int] =					0
   for i in range(len(snake_case			)			):
    UpperCamelCase_, UpperCamelCase_					:						Union[str, Any] =					task_id_and_num_shards[i]
    for shard_id in range(snake_case			):
     args.append([task_id, shard_id, global_shard_id]			)
     global_shard_id += 1
   self._spark.sparkContext.parallelize(snake_case   ,      len(snake_case			)			).map(lambda snake_case			: _rename_shard(*snake_case			)			).collect()
  else:
   # don't use any pattern
   UpperCamelCase_					:						Tuple =					0
   UpperCamelCase_					:						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(snake_case   ,      ''			)   ,      )


 def 						SCREAMING_SNAKE_CASE__   ( self	:							Tuple   ,      snake_case	:							"datasets.SplitGenerator"   ,      )   ->						SparkExamplesIterable:

  """simple docstring"""





  return SparkExamplesIterable(self.df			)



