
"""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


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

if TYPE_CHECKING:
				import pyspark




@dataclass
class 			_SCREAMING_SNAKE_CASE					(			datasets.BuilderConfig							):
			UpperCAmelCase_  :Optional[datasets.Features]				 =						None






def 						_snake_case       (	lowercase__		:  "pyspark.sql.DataFrame"      ,     lowercase__		:  List[int]      ,     )  ->				Any:


			'''simple docstring'''





			import pyspark

			def generate_fn():
						lowerCAmelCase_    :List[Any]			     =  df.select("""*"""      ,     pyspark.sql.functions.spark_partition_id().alias("""part_id"""		)		)
						for partition_id in partition_order:
									lowerCAmelCase_    :Optional[int]			     =  df_with_partition_id.select("""*"""		).where(f"""part_id = {partition_id}"""		).drop("""part_id"""		)
									lowerCAmelCase_    :Optional[Any]			     =  partition_df.collect()
									lowerCAmelCase_    :Dict			     =  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       ,   __A       ,   __A=None       ,   )			->     Optional[Any]:
						lowerCAmelCase_    :List[str]			     =  df
						lowerCAmelCase_    :str			     =  partition_order or range(self.df.rdd.getNumPartitions()       )
						lowerCAmelCase_    :int			     =  _generate_iterable_examples(self.df       ,   self.partition_order       )






			def __iter__(					self       )			->     Tuple:
						yield from self.generate_examples_fn()






			def  __lowerCAmelCase (					self       ,   __A       )			->     "SparkExamplesIterable":
						lowerCAmelCase_    :List[Any]			     =  list(range(self.df.rdd.getNumPartitions()       )       )
						generator.shuffle(__A       )
						return SparkExamplesIterable(self.df       ,   partition_order=__A       )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       )			->     "SparkExamplesIterable":
						lowerCAmelCase_    :Optional[Any]			     =  self.split_shard_indices_by_worker(__A       ,   __A       )
						return SparkExamplesIterable(self.df       ,   partition_order=__A       )






			@property
			def  __lowerCAmelCase (					self       )			->     int:
						return len(self.partition_order       )




class 			_SCREAMING_SNAKE_CASE					(			datasets.DatasetBuilder							):
			UpperCAmelCase_  :Optional[Any]				 =						SparkConfig






			def __init__(					self       ,   __A       ,   __A = None       ,   __A = None       ,   **__A       ,   )			->     int:
						import pyspark

						lowerCAmelCase_    :Tuple			     =  pyspark.sql.SparkSession.builder.getOrCreate()
						lowerCAmelCase_    :Union[str, Any]			     =  df
						lowerCAmelCase_    :Optional[Any]			     =  working_dir

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






			def  __lowerCAmelCase (					self       )			->     int:
						# Returns the path of the created file.
						def create_cache_and_write_probe(__A       ):
									# 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       )
									lowerCAmelCase_    :Union[str, 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:
									lowerCAmelCase_    :int			     =  (
									    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  __lowerCAmelCase (					self       )			->     Optional[Any]:
						return datasets.DatasetInfo(features=self.config.features       )






			def  __lowerCAmelCase (					self       ,   __A       )			->     Any:
						return [datasets.SplitGenerator(name=datasets.Split.TRAIN       )]






			def  __lowerCAmelCase (					self       ,   __A       )			->     Union[str, Any]:
						import pyspark

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

						lowerCAmelCase_    :Tuple			     =  self.df.count()
						lowerCAmelCase_    :Union[str, 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.
						lowerCAmelCase_    :Tuple			     =  (
						    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
						)
						lowerCAmelCase_    :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.
									lowerCAmelCase_    :str			     =  min(__A       ,   int(approx_total_size / max_shard_size       )       )
									lowerCAmelCase_    :Optional[int]			     =  self.df.repartition(__A       )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   )			->     Iterable[Tuple[int, bool, Union[int, tuple]]]:
						import pyspark

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

						def write_arrow(__A       ):
									# Within the same SparkContext, no two task attempts will share the same attempt ID.
									lowerCAmelCase_    :Dict			     =  pyspark.TaskContext().taskAttemptId()
									lowerCAmelCase_    :int			     =  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"""]       ,   )
									lowerCAmelCase_    :Tuple			     =  0
									lowerCAmelCase_    :List[str]			     =  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       ,   )
									lowerCAmelCase_    :int			     =  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:
															lowerCAmelCase_     ,       lowerCAmelCase_    :int			     =  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_    :int			     =  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       ,   )
												lowerCAmelCase_    :Any			     =  pa.Table.from_batches([batch]       )
												writer.write_table(__A       )

									if writer._num_bytes > 0:
												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"""]       ,   )

									if working_fpath != fpath:
												for file in os.listdir(os.path.dirname(__A       )       ):
															lowerCAmelCase_    :Optional[int]			     =  os.path.join(os.path.dirname(__A       )       ,   os.path.basename(__A       )       )
															shutil.move(__A       ,   __A       )

						lowerCAmelCase_    :Optional[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  __lowerCAmelCase (					self       ,   __A       ,   __A = "arrow"       ,   __A = None       ,   __A = None       ,   **__A       ,   )			->     Any:
						self._validate_cache_dir()

						lowerCAmelCase_    :Tuple			     =  convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE       )
						self._repartition_df_if_needed(__A       )
						lowerCAmelCase_    :Optional[Any]			     =  not is_remote_filesystem(self._fs       )
						lowerCAmelCase_    :Tuple			     =  os.path.join if is_local else posixpath.join

						lowerCAmelCase_    :List[Any]			     =  """-TTTTT-SSSSS-of-NNNNN"""
						lowerCAmelCase_    :int			     =  f"""{self.name}-{split_generator.name}{SUFFIX}.{file_format}"""
						lowerCAmelCase_    :Optional[Any]			     =  path_join(self._output_dir       ,   __A       )

						lowerCAmelCase_    :Dict			     =  0
						lowerCAmelCase_    :Any			     =  0
						lowerCAmelCase_    :str			     =  0
						lowerCAmelCase_    :Union[str, Any]			     =  []
						lowerCAmelCase_    :List[str]			     =  []

						for task_id, content in self._prepare_split_single(__A       ,   __A       ,   __A       ):
									(
									    (
									    lowerCAmelCase_
									)     ,       (
									    lowerCAmelCase_
									)     ,       (
									    lowerCAmelCase_
									)     ,       (
									    lowerCAmelCase_
									)     ,       
									)    :List[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(__A       )

						lowerCAmelCase_    :Optional[int]			     =  total_num_examples
						lowerCAmelCase_    :Tuple			     =  total_num_bytes

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

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

									lowerCAmelCase_    :Tuple			     =  []
									lowerCAmelCase_    :Tuple			     =  0
									for i in range(len(__A       )       ):
												lowerCAmelCase_     ,       lowerCAmelCase_    :Dict			     =  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
									lowerCAmelCase_    :Optional[int]			     =  0
									lowerCAmelCase_    :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  __lowerCAmelCase (					self       ,   __A       ,   )			->     SparkExamplesIterable:
						return SparkExamplesIterable(self.df       )













"""simple docstring"""





import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


SCREAMING_SNAKE_CASE__					=   "▁"

SCREAMING_SNAKE_CASE__					=   {"vocab_file": "spiece.model"}

SCREAMING_SNAKE_CASE__					=   {
    "vocab_file": {"google/pegasus-xsum": "https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"}
}

SCREAMING_SNAKE_CASE__					=   {
    "google/pegasus-xsum": 512,
}


SCREAMING_SNAKE_CASE__					=   logging.get_logger(__name__)

class 		lowercase					(					_UpperCAmelCase					):
      _SCREAMING_SNAKE_CASE					=      VOCAB_FILES_NAMES

      _SCREAMING_SNAKE_CASE					=      VOCAB_FILES_NAMES
      _SCREAMING_SNAKE_CASE					=      PRETRAINED_VOCAB_FILES_MAP
      _SCREAMING_SNAKE_CASE					=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _SCREAMING_SNAKE_CASE					=      ['input_ids', 'attention_mask']





      def __init__(       self      ,						lowercase      ,						lowercase="<pad>"      ,						lowercase="</s>"      ,						lowercase="<unk>"      ,						lowercase="<mask_2>"      ,						lowercase="<mask_1>"      ,						lowercase=None      ,						lowercase=103      ,						lowercase = None      ,						**lowercase      ,						)		->     None:
            lowerCAmelCase  						=       offset
            if additional_special_tokens is not None:
                  if not isinstance(lowercase      ,						lowercase ):
                        raise TypeError(
                            f'additional_special_tokens should be of type {type(lowercase )}, but is'
                            f' {type(lowercase )}' )

                  lowerCAmelCase  						=       (
                      ([mask_token_sent] + additional_special_tokens)
                      if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                      else additional_special_tokens
                  )
                  # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
                  additional_special_tokens_extended += [
                      f'<unk_{i}>' for i in range(len(lowercase )      ,						self.offset - 1 )
                  ]

                  if len(set(lowercase ) ) != len(lowercase ):
                        raise ValueError(
                            """Please make sure that the provided additional_special_tokens do not contain an incorrectly"""
                            f' shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.' )
                  lowerCAmelCase  						=       additional_special_tokens_extended
            else:
                  lowerCAmelCase  						=       [mask_token_sent] if mask_token_sent is not None else []
                  additional_special_tokens += [f'<unk_{i}>' for i in range(2      ,						self.offset )]

            lowerCAmelCase  						=       {} if sp_model_kwargs is None else sp_model_kwargs

            super().__init__(
                eos_token=lowercase      ,						unk_token=lowercase      ,						mask_token=lowercase      ,						pad_token=lowercase      ,						mask_token_sent=lowercase      ,						offset=lowercase      ,						additional_special_tokens=lowercase      ,						sp_model_kwargs=self.sp_model_kwargs      ,						**lowercase      ,						)
            lowerCAmelCase  						=       mask_token_sent
            lowerCAmelCase  						=       vocab_file
            lowerCAmelCase  						=       spm.SentencePieceProcessor(**self.sp_model_kwargs )
            self.sp_model.Load(lowercase )

            # add special tokens to encoder dict
            lowerCAmelCase  						=       {
                0: self.pad_token,
                1: self.eos_token,
            }

            if self.mask_token_sent is not None:
                  self.encoder.update(
                      {
                          2: self.mask_token_sent,
                          3: self.mask_token,
                      } )

            if self.offset > 0:
                  # entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
                  # mask_token_sent is already added to list -> so start at 1
                  self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1      ,						self.offset - 1 )} )

            lowerCAmelCase  						=       {v: k for k, v in self.encoder.items()}





      @property
      def     _snake_case			(       self )		->     int:
            return len(self.sp_model ) + self.offset





      def     _snake_case			(       self )		->     Dict[str, int]:
            lowerCAmelCase  						=       {self.convert_ids_to_tokens(lowercase ): i for i in range(self.vocab_size )}
            vocab.update(self.added_tokens_encoder )
            return vocab





      def __getstate__(       self )		->     Optional[int]:
            lowerCAmelCase  						=       self.__dict__.copy()
            lowerCAmelCase  						=       None
            return state





      def __setstate__(       self      ,						lowercase )		->     List[Any]:
            lowerCAmelCase  						=       d

            # for backward compatibility
            if not hasattr(self      ,						"""sp_model_kwargs""" ):
                  lowerCAmelCase  						=       {}

            lowerCAmelCase  						=       spm.SentencePieceProcessor(**self.sp_model_kwargs )
            self.sp_model.Load(self.vocab_file )





      def     _snake_case			(       self      ,						lowercase )		->     List[str]:
            return self.sp_model.encode(lowercase      ,						out_type=lowercase )





      def     _snake_case			(       self      ,						lowercase )		->     int:
            if token in self.decoder:
                  return self.decoder[token]
            elif token in self.added_tokens_decoder:
                  return self.added_tokens_decoder[token]
            lowerCAmelCase  						=       self.sp_model.piece_to_id(lowercase )
            return sp_id + self.offset





      def     _snake_case			(       self      ,						lowercase )		->     str:
            if index in self.encoder:
                  return self.encoder[index]
            elif index in self.added_tokens_encoder:
                  return self.added_tokens_encoder[index]
            else:
                  lowerCAmelCase  						=       self.sp_model.IdToPiece(index - self.offset )
            return token





      def     _snake_case			(       self      ,						lowercase )		->     Optional[int]:
            lowerCAmelCase  						=       []
            lowerCAmelCase  						=       """"""
            for token in tokens:
                  # make sure that special tokens are not decoded using sentencepiece model
                  if token in self.all_special_tokens:
                        out_string += self.sp_model.decode(lowercase ) + token
                        lowerCAmelCase  						=       []
                  else:
                        current_sub_tokens.append(lowercase )
            out_string += self.sp_model.decode(lowercase )
            return out_string.strip()





      def     _snake_case			(       self      ,						lowercase=False )		->     Tuple:
            return 1





      def     _snake_case			(       self      ,						lowercase )		->     Tuple:
            lowerCAmelCase  						=       set(self.all_special_ids )  # call it once instead of inside list comp
            all_special_ids.remove(self.unk_token_id )  # <unk> is only sometimes special

            return [1 if x in all_special_ids else 0 for x in seq]





      def     _snake_case			(       self      ,						lowercase      ,						lowercase = None      ,						lowercase = False )		->     List[int]:
            if already_has_special_tokens:
                  return self._special_token_mask(lowercase )
            elif token_ids_a is None:
                  return self._special_token_mask(lowercase ) + [1]
            else:
                  return self._special_token_mask(token_ids_a + token_ids_a ) + [1]





      def     _snake_case			(       self      ,						lowercase      ,						lowercase=None )		->     List[int]:
            if token_ids_a is None:
                  return token_ids_a + [self.eos_token_id]
            # We don't expect to process pairs, but leave the pair logic for API consistency
            return token_ids_a + token_ids_a + [self.eos_token_id]





      def     _snake_case			(       self      ,						lowercase      ,						lowercase = None )		->     Tuple[str]:
            if not os.path.isdir(lowercase ):
                  logger.error(f'Vocabulary path ({save_directory}) should be a directory' )
                  return
            lowerCAmelCase  						=       os.path.join(
                lowercase      ,						(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""] )

            if os.path.abspath(self.vocab_file ) != os.path.abspath(lowercase ) and os.path.isfile(self.vocab_file ):
                  copyfile(self.vocab_file      ,						lowercase )
            elif not os.path.isfile(self.vocab_file ):
                  with open(lowercase      ,						"""wb""" ) as fi:
                        lowerCAmelCase  						=       self.sp_model.serialized_model_proto()
                        fi.write(lowercase )

            return (out_vocab_file,)






