




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


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

if TYPE_CHECKING:
			import pyspark






@dataclass
class lowerCAmelCase__	(     datasets.BuilderConfig						):






			'''simple docstring'''



			lowerCamelCase__      =			None







def      _snake_case			(				lowercase__							,    lowercase__							,    ):
				import pyspark

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

				return generate_fn






class lowerCAmelCase__	(     _BaseExamplesIterable						):






			'''simple docstring'''



			def __init__(  self  ,      lowercase  ,      lowercase=None  ,      ):
							_lowerCamelCase     :			Any    =    df
							_lowerCamelCase     :			List[Any]    =    partition_order or range(self.df.rdd.getNumPartitions()  )
							_lowerCamelCase     :			Union[str, Any]    =    _generate_iterable_examples(self.df  ,      self.partition_order  )
			def __iter__(  self  ):
							yield from self.generate_examples_fn()
			def 					A_    (  self  ,      lowercase  ):
							_lowerCamelCase     :			List[str]    =    list(range(self.df.rdd.getNumPartitions()  )  )
							generator.shuffle(lowercase  )
							return SparkExamplesIterable(self.df  ,      partition_order=lowercase  )
			def 					A_    (  self  ,      lowercase  ,      lowercase  ):
							_lowerCamelCase     :			str    =    self.split_shard_indices_by_worker(lowercase  ,      lowercase  )
							return SparkExamplesIterable(self.df  ,      partition_order=lowercase  )






			@property
			def 					A_    (  self  ):
							return len(self.partition_order  )






class lowerCAmelCase__	(     datasets.DatasetBuilder						):






			'''simple docstring'''



			lowerCamelCase__      =			SparkConfig
			def __init__(  self  ,      lowercase  ,      lowercase = None  ,      lowercase = None  ,      **lowercase  ,      ):
							import pyspark

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

							super().__init__(
							    cache_dir=lowercase  ,      config_name=str(self.df.semanticHash()  )  ,      **lowercase  ,      )
			def 					A_    (  self  ):
							# Returns the path of the created file.
							def create_cache_and_write_probe(lowercase  ):
											# 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=lowercase  )
											_lowerCamelCase     :			List[str]    =    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(lowercase  ,      '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(lowercase  ).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 					A_    (  self  ):
							return datasets.DatasetInfo(features=self.config.features  )
			def 					A_    (  self  ,      lowercase  ):
							return [datasets.SplitGenerator(name=datasets.Split.TRAIN  )]
			def 					A_    (  self  ,      lowercase  ):
							import pyspark

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

							_lowerCamelCase     :			int    =    self.df.count()
							_lowerCamelCase     :			Dict    =    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     :			Optional[int]    =    (
							    self.df.limit(lowercase  )
							    .repartition(1  )
							    .mapInArrow(lowercase  ,      'batch_bytes: long'  )
							    .agg(pyspark.sql.functions.sum('batch_bytes'  ).alias('sample_bytes'  )  )
							    .collect()[0]
							    .sample_bytes
							    / sample_num_rows
							)
							_lowerCamelCase     :			Tuple    =    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     :			Dict    =    min(lowercase  ,      int(approx_total_size / max_shard_size  )  )
											_lowerCamelCase     :			List[str]    =    self.df.repartition(lowercase  )
			def 					A_    (  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      ):
							import pyspark

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

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

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

							_lowerCamelCase     :			Any    =    (
							    self.df.mapInArrow(lowercase  ,      '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 					A_    (  self  ,      lowercase  ,      lowercase = "arrow"  ,      lowercase = None  ,      lowercase = None  ,      **lowercase  ,      ):
							self._validate_cache_dir()

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

							_lowerCamelCase     :			Optional[Any]    =    '-TTTTT-SSSSS-of-NNNNN'
							_lowerCamelCase     :			Any    =    F'''{self.name}-{split_generator.name}{SUFFIX}.{file_format}'''
							_lowerCamelCase     :			Any    =    path_join(self._output_dir  ,      lowercase  )

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

							for task_id, content in self._prepare_split_single(lowercase  ,      lowercase  ,      lowercase  ):
											(
											    (
											    _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(lowercase  )

							_lowerCamelCase     :			int    =    total_num_examples
							_lowerCamelCase     :			Optional[int]    =    total_num_bytes

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

											# use the -SSSSS-of-NNNNN pattern
											def _rename_shard(
											    lowercase  ,      lowercase  ,      lowercase  ,      ):
															rename(
															    lowercase  ,      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     :			List[Any]    =    []
											_lowerCamelCase     :			Union[str, Any]    =    0
											for i in range(len(lowercase  )  ):
															_lowerCamelCase, _lowerCamelCase     :			List[str]    =    task_id_and_num_shards[i]
															for shard_id in range(lowercase  ):
																			args.append([task_id, shard_id, global_shard_id]  )
																			global_shard_id += 1
											self._spark.sparkContext.parallelize(lowercase  ,      len(lowercase  )  ).map(lambda lowercase  : _rename_shard(*lowercase  )  ).collect()
							else:
											# don't use any pattern
											_lowerCamelCase     :			Dict    =    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(lowercase  ,      ''  )  ,      )






			def 					A_    (  self  ,      lowercase  ,      ):
							return SparkExamplesIterable(self.df  )





"""simple docstring"""






import os
from typing import List, Optional, Union

from ...image_processing_utils import BatchFeature
from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType
from ..auto import AutoTokenizer






class lowerCAmelCase__	(     lowercase						):






			'''simple docstring'''



			lowerCamelCase__      =			["""image_processor""", """tokenizer"""]
			lowerCamelCase__      =			"""BlipImageProcessor"""
			lowerCamelCase__      =			"""AutoTokenizer"""
			def __init__(  self  ,      lowercase  ,      lowercase  ,      lowercase  ):
							super().__init__(lowercase  ,      lowercase  )

							# add QFormer tokenizer
							_lowerCamelCase     :			int    =    qformer_tokenizer
			def __call__(  self  ,      lowercase = None  ,      lowercase = None  ,      lowercase = True  ,      lowercase = False  ,      lowercase = None  ,      lowercase = None  ,      lowercase = 0  ,      lowercase = None  ,      lowercase = None  ,      lowercase = False  ,      lowercase = False  ,      lowercase = False  ,      lowercase = False  ,      lowercase = False  ,      lowercase = True  ,      lowercase = None  ,      **lowercase  ,      ):
							if images is None and text is None:
											raise ValueError('You have to specify at least images or text.'  )

							_lowerCamelCase     :			int    =    BatchFeature()

							if text is not None:
											_lowerCamelCase     :			List[str]    =    self.tokenizer(
											    text=lowercase  ,      add_special_tokens=lowercase  ,      padding=lowercase  ,      truncation=lowercase  ,      max_length=lowercase  ,      stride=lowercase  ,      pad_to_multiple_of=lowercase  ,      return_attention_mask=lowercase  ,      return_overflowing_tokens=lowercase  ,      return_special_tokens_mask=lowercase  ,      return_offsets_mapping=lowercase  ,      return_token_type_ids=lowercase  ,      return_length=lowercase  ,      verbose=lowercase  ,      return_tensors=lowercase  ,      **lowercase  ,      )
											encoding.update(lowercase  )
											_lowerCamelCase     :			List[str]    =    self.qformer_tokenizer(
											    text=lowercase  ,      add_special_tokens=lowercase  ,      padding=lowercase  ,      truncation=lowercase  ,      max_length=lowercase  ,      stride=lowercase  ,      pad_to_multiple_of=lowercase  ,      return_attention_mask=lowercase  ,      return_overflowing_tokens=lowercase  ,      return_special_tokens_mask=lowercase  ,      return_offsets_mapping=lowercase  ,      return_token_type_ids=lowercase  ,      return_length=lowercase  ,      verbose=lowercase  ,      return_tensors=lowercase  ,      **lowercase  ,      )
											_lowerCamelCase     :			List[Any]    =    qformer_text_encoding.pop('input_ids'  )
											_lowerCamelCase     :			Tuple    =    qformer_text_encoding.pop('attention_mask'  )

							if images is not None:
											_lowerCamelCase     :			int    =    self.image_processor(lowercase  ,      return_tensors=lowercase  )
											encoding.update(lowercase  )

							return encoding
			def 					A_    (  self  ,      *lowercase  ,      **lowercase  ):
							return self.tokenizer.batch_decode(*lowercase  ,      **lowercase  )
			def 					A_    (  self  ,      *lowercase  ,      **lowercase  ):
							return self.tokenizer.decode(*lowercase  ,      **lowercase  )
			@property
			# Copied from transformers.models.blip.processing_blip.BlipProcessor.model_input_names
			def 					A_    (  self  ):
							_lowerCamelCase     :			Union[str, Any]    =    self.tokenizer.model_input_names
							_lowerCamelCase     :			Any    =    self.image_processor.model_input_names
							return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names  )  )
			def 					A_    (  self  ,      lowercase  ,      **lowercase  ):
							if os.path.isfile(lowercase  ):
											raise ValueError(F'''Provided path ({save_directory}) should be a directory, not a file'''  )
							os.makedirs(lowercase  ,      exist_ok=lowercase  )
							_lowerCamelCase     :			Optional[Any]    =    os.path.join(lowercase  ,      'qformer_tokenizer'  )
							self.qformer_tokenizer.save_pretrained(lowercase  )
							return super().save_pretrained(lowercase  ,      **lowercase  )






			@classmethod
			def 					A_    (  cls  ,      lowercase  ,      **lowercase  ):
							_lowerCamelCase     :			Union[str, Any]    =    AutoTokenizer.from_pretrained(lowercase  ,      subfolder='qformer_tokenizer'  )
							_lowerCamelCase     :			Dict    =    cls._get_arguments_from_pretrained(lowercase  ,      **lowercase  )
							args.append(lowercase  )
							return cls(*lowercase  )