


'''simple docstring'''







import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader


class  lowerCAmelCase_						(       snake_case__      ):





  """simple docstring"""



  def __init__(					self     :							List[str]				,							SCREAMING_SNAKE_CASE__     :							NestedDataStructureLike[PathLike]				,							SCREAMING_SNAKE_CASE__     :							Optional[NamedSplit] = None				,							SCREAMING_SNAKE_CASE__     :							Optional[Features] = None				,							SCREAMING_SNAKE_CASE__     :							str = None				,							SCREAMING_SNAKE_CASE__     :							bool = False				,							SCREAMING_SNAKE_CASE__     :							bool = False				,							SCREAMING_SNAKE_CASE__     :							Optional[str] = None				,							SCREAMING_SNAKE_CASE__     :							Optional[int] = None				,							**SCREAMING_SNAKE_CASE__     :							Optional[int]				,							):





       '''simple docstring'''

       super().__init__(
           SCREAMING_SNAKE_CASE__				,							split=SCREAMING_SNAKE_CASE__				,							features=SCREAMING_SNAKE_CASE__				,							cache_dir=SCREAMING_SNAKE_CASE__				,							keep_in_memory=SCREAMING_SNAKE_CASE__				,							streaming=SCREAMING_SNAKE_CASE__				,							num_proc=SCREAMING_SNAKE_CASE__				,							**SCREAMING_SNAKE_CASE__				,							)
       __a       =   field
       __a       =   path_or_paths if isinstance(SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						) else {self.split: path_or_paths}
       __a       =   Json(
           cache_dir=SCREAMING_SNAKE_CASE__				,							data_files=SCREAMING_SNAKE_CASE__				,							features=SCREAMING_SNAKE_CASE__				,							field=SCREAMING_SNAKE_CASE__				,							**SCREAMING_SNAKE_CASE__				,							)


  def __a						(					self     :							str						):





       '''simple docstring'''

       if self.streaming:
            __a       =   self.builder.as_streaming_dataset(split=self.split						)
       # Build regular (map-style) dataset
       else:
            __a       =   None
            __a       =   None
            __a       =   None
            __a       =   None

            self.builder.download_and_prepare(
                download_config=SCREAMING_SNAKE_CASE__				,							download_mode=SCREAMING_SNAKE_CASE__				,							verification_mode=SCREAMING_SNAKE_CASE__				,							base_path=SCREAMING_SNAKE_CASE__				,							num_proc=self.num_proc				,							)
            __a       =   self.builder.as_dataset(
                split=self.split				,							verification_mode=SCREAMING_SNAKE_CASE__				,							in_memory=self.keep_in_memory						)
       return dataset


class  lowerCAmelCase_						:





  """simple docstring"""



  def __init__(					self     :							Optional[Any]				,							SCREAMING_SNAKE_CASE__     :							Dataset				,							SCREAMING_SNAKE_CASE__     :							Union[PathLike, BinaryIO]				,							SCREAMING_SNAKE_CASE__     :							Optional[int] = None				,							SCREAMING_SNAKE_CASE__     :							Optional[int] = None				,							**SCREAMING_SNAKE_CASE__     :							Optional[int]				,							):





       '''simple docstring'''

       if num_proc is not None and num_proc <= 0:
            raise ValueError(f'''num_proc {num_proc} must be an integer > 0.'''						)

       __a       =   dataset
       __a       =   path_or_buf
       __a       =   batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
       __a       =   num_proc
       __a       =   """utf-8"""
       __a       =   to_json_kwargs

  def __a						(					self     :							Optional[int]						):





       '''simple docstring'''

       __a       =   self.to_json_kwargs.pop("""path_or_buf"""				,							SCREAMING_SNAKE_CASE__						)
       __a       =   self.to_json_kwargs.pop("""orient"""				,							"""records"""						)
       __a       =   self.to_json_kwargs.pop("""lines"""				,							True if orient == """records""" else False						)
       __a       =   self.to_json_kwargs.pop("""index"""				,							False if orient in ["""split""", """table"""] else True						)
       __a       =   self.to_json_kwargs.pop("""compression"""				,							SCREAMING_SNAKE_CASE__						)

       if compression not in [None, "infer", "gzip", "bz2", "xz"]:
            raise NotImplementedError(f'''`datasets` currently does not support {compression} compression'''						)

       if isinstance(self.path_or_buf				,							(str, bytes, os.PathLike)						):
            with fsspec.open(self.path_or_buf				,							"""wb"""				,							compression=SCREAMING_SNAKE_CASE__						) as buffer:
                 __a       =   self._write(file_obj=SCREAMING_SNAKE_CASE__				,							orient=SCREAMING_SNAKE_CASE__				,							lines=SCREAMING_SNAKE_CASE__				,							index=SCREAMING_SNAKE_CASE__				,							**self.to_json_kwargs						)
       else:
            if compression:
                 raise NotImplementedError(
                     f'''The compression parameter is not supported when writing to a buffer, but compression={compression}'''
                     """ was passed. Please provide a local path instead."""						)
            __a       =   self._write(
                file_obj=self.path_or_buf				,							orient=SCREAMING_SNAKE_CASE__				,							lines=SCREAMING_SNAKE_CASE__				,							index=SCREAMING_SNAKE_CASE__				,							**self.to_json_kwargs						)
       return written

  def __a						(					self     :							Union[str, Any]				,							SCREAMING_SNAKE_CASE__     :							List[str]						):





       '''simple docstring'''

       __a ,						__a ,						__a ,						__a ,						__a       =   args

       __a       =   query_table(
           table=self.dataset.data				,							key=slice(SCREAMING_SNAKE_CASE__				,							offset + self.batch_size						)				,							indices=self.dataset._indices				,							)
       __a       =   batch.to_pandas().to_json(
           path_or_buf=SCREAMING_SNAKE_CASE__				,							orient=SCREAMING_SNAKE_CASE__				,							lines=SCREAMING_SNAKE_CASE__				,							index=SCREAMING_SNAKE_CASE__				,							**SCREAMING_SNAKE_CASE__						)
       if not json_str.endswith("""\n"""						):
            json_str += "\n"
       return json_str.encode(self.encoding						)


  def __a						(					self     :							List[Any]				,							SCREAMING_SNAKE_CASE__     :							BinaryIO				,							SCREAMING_SNAKE_CASE__     :							Dict				,							SCREAMING_SNAKE_CASE__     :							Dict				,							SCREAMING_SNAKE_CASE__     :							List[Any]				,							**SCREAMING_SNAKE_CASE__     :							Dict				,							):





       '''simple docstring'''

       __a       =   0

       if self.num_proc is None or self.num_proc == 1:
            for offset in logging.tqdm(
                range(0				,							len(self.dataset						)				,							self.batch_size						)				,							unit="""ba"""				,							disable=not logging.is_progress_bar_enabled()				,							desc="""Creating json from Arrow format"""				,							):
                 __a       =   self._batch_json((offset, orient, lines, index, to_json_kwargs)						)
                 written += file_obj.write(SCREAMING_SNAKE_CASE__						)
       else:
            __a ,						__a       =   len(self.dataset						), self.batch_size
            with multiprocessing.Pool(self.num_proc						) as pool:
                 for json_str in logging.tqdm(
                     pool.imap(
                         self._batch_json				,							[(offset, orient, lines, index, to_json_kwargs) for offset in range(0				,							SCREAMING_SNAKE_CASE__				,							SCREAMING_SNAKE_CASE__						)]				,							)				,							total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size				,							unit="""ba"""				,							disable=not logging.is_progress_bar_enabled()				,							desc="""Creating json from Arrow format"""				,							):
                      written += file_obj.write(SCREAMING_SNAKE_CASE__						)

       return written










'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


SCREAMING_SNAKE_CASE_      =			{
    'configuration_rag': ['RagConfig'],
    'retrieval_rag': ['RagRetriever'],
    'tokenization_rag': ['RagTokenizer'],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE_      =			[
     'RagModel',
     'RagPreTrainedModel',
     'RagSequenceForGeneration',
     'RagTokenForGeneration',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE_      =			[
     'TFRagModel',
     'TFRagPreTrainedModel',
     'TFRagSequenceForGeneration',
     'TFRagTokenForGeneration',
 ]


if TYPE_CHECKING:
 from .configuration_rag import RagConfig
 from .retrieval_rag import RagRetriever
 from .tokenization_rag import RagTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_rag import (
      TFRagModel,
      TFRagPreTrainedModel,
      TFRagSequenceForGeneration,
      TFRagTokenForGeneration,
  )

else:
 import sys

 SCREAMING_SNAKE_CASE_      =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






