

'''simple docstring'''




import json
import os
import shutil
import tempfile
from unittest import TestCase

from transformers import BartTokenizer, BartTokenizerFast, DPRQuestionEncoderTokenizer, DPRQuestionEncoderTokenizerFast
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_tokenizers, require_torch, slow
from transformers.utils import is_datasets_available, is_faiss_available, is_torch_available


if is_torch_available() and is_datasets_available() and is_faiss_available():
  from transformers.models.rag.configuration_rag import RagConfig
  from transformers.models.rag.tokenization_rag import RagTokenizer



@require_faiss
@require_torch
class 							lowercase_						(		_lowerCamelCase			):

  '''simple docstring'''





  def     snake_case_      (  self  )				->					Any:


    """simple docstring"""
    UpperCAmelCase    				=						tempfile.mkdtemp()
    UpperCAmelCase    				=						8

    # DPR tok
    UpperCAmelCase    				=						[
        '[UNK]',
        '[CLS]',
        '[SEP]',
        '[PAD]',
        '[MASK]',
        'want',
        '##want',
        '##ed',
        'wa',
        'un',
        'runn',
        '##ing',
        ',',
        'low',
        'lowest',
    ]
    UpperCAmelCase    				=						os.path.join(self.tmpdirname ,			'dpr_tokenizer'  )
    os.makedirs(lowercase_ ,			exist_ok=lowercase_  )
    UpperCAmelCase    				=						os.path.join(lowercase_ ,			DPR_VOCAB_FILES_NAMES['vocab_file']  )
    with open(self.vocab_file ,			'w' ,			encoding='utf-8'  ) as vocab_writer:
      vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]  )  )

    # BART tok
    UpperCAmelCase    				=						[
        'l',
        'o',
        'w',
        'e',
        'r',
        's',
        't',
        'i',
        'd',
        'n',
        '\u0120',
        '\u0120l',
        '\u0120n',
        '\u0120lo',
        '\u0120low',
        'er',
        '\u0120lowest',
        '\u0120newer',
        '\u0120wider',
        '<unk>',
    ]
    UpperCAmelCase    				=						dict(zip(lowercase_ ,			range(len(lowercase_  )  )  )  )
    UpperCAmelCase    				=						['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
    UpperCAmelCase    				=						{'unk_token': '<unk>'}

    UpperCAmelCase    				=						os.path.join(self.tmpdirname ,			'bart_tokenizer'  )
    os.makedirs(lowercase_ ,			exist_ok=lowercase_  )
    UpperCAmelCase    				=						os.path.join(lowercase_ ,			BART_VOCAB_FILES_NAMES['vocab_file']  )
    UpperCAmelCase    				=						os.path.join(lowercase_ ,			BART_VOCAB_FILES_NAMES['merges_file']  )
    with open(self.vocab_file ,			'w' ,			encoding='utf-8'  ) as fp:
      fp.write(json.dumps(lowercase_  ) + '\n'  )
    with open(self.merges_file ,			'w' ,			encoding='utf-8'  ) as fp:
      fp.write('\n'.join(lowercase_  )  )
  def     snake_case_      (  self  )				->					Any:


    """simple docstring"""
    return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,			'dpr_tokenizer'  )  )
  def     snake_case_      (  self  )				->					Dict:


    """simple docstring"""
    return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname ,			'bart_tokenizer'  )  )
  def     snake_case_      (  self  )				->					Any:


    """simple docstring"""
    shutil.rmtree(self.tmpdirname  )
  @require_tokenizers
  def     snake_case_      (  self  )				->					Optional[int]:


    """simple docstring"""
    UpperCAmelCase    				=						os.path.join(self.tmpdirname ,			'rag_tokenizer'  )
    UpperCAmelCase    				=						RagConfig(question_encoder=DPRConfig().to_dict() ,			generator=BartConfig().to_dict()  )
    UpperCAmelCase    				=						RagTokenizer(question_encoder=self.get_dpr_tokenizer() ,			generator=self.get_bart_tokenizer()  )
    rag_config.save_pretrained(lowercase_  )
    rag_tokenizer.save_pretrained(lowercase_  )
    UpperCAmelCase    				=						RagTokenizer.from_pretrained(lowercase_ ,			config=lowercase_  )
    self.assertIsInstance(new_rag_tokenizer.question_encoder ,			lowercase_  )
    self.assertEqual(new_rag_tokenizer.question_encoder.get_vocab() ,			rag_tokenizer.question_encoder.get_vocab()  )
    self.assertIsInstance(new_rag_tokenizer.generator ,			lowercase_  )
    self.assertEqual(new_rag_tokenizer.generator.get_vocab() ,			rag_tokenizer.generator.get_vocab()  )
  @slow
  def     snake_case_      (  self  )				->					List[str]:


    """simple docstring"""
    UpperCAmelCase    				=						RagTokenizer.from_pretrained('facebook/rag-token-nq'  )
    UpperCAmelCase    				=						[
        'who got the first nobel prize in physics',
        'when is the next deadpool movie being released',
        'which mode is used for short wave broadcast service',
        'who is the owner of reading football club',
        'when is the next scandal episode coming out',
        'when is the last time the philadelphia won the superbowl',
        'what is the most current adobe flash player version',
        'how many episodes are there in dragon ball z',
        'what is the first step in the evolution of the eye',
        'where is gall bladder situated in human body',
        'what is the main mineral in lithium batteries',
        'who is the president of usa right now',
        'where do the greasers live in the outsiders',
        'panda is a national animal of which country',
        'what is the name of manchester united stadium',
    ]
    UpperCAmelCase    				=						tokenizer(lowercase_  )
    self.assertIsNotNone(lowercase_  )





  @slow
  def     snake_case_      (  self  )				->					int:


    """simple docstring"""
    UpperCAmelCase    				=						RagTokenizer.from_pretrained('facebook/rag-sequence-nq'  )
    UpperCAmelCase    				=						[
        'who got the first nobel prize in physics',
        'when is the next deadpool movie being released',
        'which mode is used for short wave broadcast service',
        'who is the owner of reading football club',
        'when is the next scandal episode coming out',
        'when is the last time the philadelphia won the superbowl',
        'what is the most current adobe flash player version',
        'how many episodes are there in dragon ball z',
        'what is the first step in the evolution of the eye',
        'where is gall bladder situated in human body',
        'what is the main mineral in lithium batteries',
        'who is the president of usa right now',
        'where do the greasers live in the outsiders',
        'panda is a national animal of which country',
        'what is the name of manchester united stadium',
    ]
    UpperCAmelCase    				=						tokenizer(lowercase_  )
    self.assertIsNotNone(lowercase_  )








'''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 		a							(  _lowerCamelCase	):


				def __init__(						self       :   Optional[int]			,							lowercase_       :   NestedDataStructureLike[PathLike]			,							lowercase_       :   Optional[NamedSplit] = None			,							lowercase_       :   Optional[Features] = None			,							lowercase_       :   str = None			,							lowercase_       :   bool = False			,							lowercase_       :   bool = False			,							lowercase_       :   Optional[str] = None			,							lowercase_       :   Optional[int] = None			,							**lowercase_       :   Optional[Any]			,							):
									super().__init__(
									    lowercase_			,							split=lowercase_			,							features=lowercase_			,							cache_dir=lowercase_			,							keep_in_memory=lowercase_			,							streaming=lowercase_			,							num_proc=lowercase_			,							**lowercase_			,							)
									snake_case_		     =							field
									snake_case_		     =							path_or_paths if isinstance(lowercase_			,							lowercase_   ) else {self.split: path_or_paths}
									snake_case_		     =							Json(
									    cache_dir=lowercase_			,							data_files=lowercase_			,							features=lowercase_			,							field=lowercase_			,							**lowercase_			,							)


				def     A_					(						self       :   Optional[Any]   ):
									# Build iterable dataset
									if self.streaming:
														snake_case_		     =							self.builder.as_streaming_dataset(split=self.split   )
									# Build regular (map-style) dataset
									else:
														snake_case_		     =							None
														snake_case_		     =							None
														snake_case_		     =							None
														snake_case_		     =							None

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







class 		a							:


				def __init__(						self       :   Dict			,							lowercase_       :   Dataset			,							lowercase_       :   Union[PathLike, BinaryIO]			,							lowercase_       :   Optional[int] = None			,							lowercase_       :   Optional[int] = None			,							**lowercase_       :   List[str]			,							):
									if num_proc is not None and num_proc <= 0:
														raise ValueError(F"num_proc {num_proc} must be an integer > 0."   )

									snake_case_		     =							dataset
									snake_case_		     =							path_or_buf
									snake_case_		     =							batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
									snake_case_		     =							num_proc
									snake_case_		     =							'''utf-8'''
									snake_case_		     =							to_json_kwargs


				def     A_					(						self       :   str   ):
									snake_case_		     =							self.to_json_kwargs.pop('''path_or_buf'''			,							lowercase_   )
									snake_case_		     =							self.to_json_kwargs.pop('''orient'''			,							'''records'''   )
									snake_case_		     =							self.to_json_kwargs.pop('''lines'''			,							True if orient == '''records''' else False   )
									snake_case_		     =							self.to_json_kwargs.pop('''index'''			,							False if orient in ['''split''', '''table'''] else True   )
									snake_case_		     =							self.to_json_kwargs.pop('''compression'''			,							lowercase_   )

									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=lowercase_   ) as buffer:
																			snake_case_		     =							self._write(file_obj=lowercase_			,							orient=lowercase_			,							lines=lowercase_			,							index=lowercase_			,							**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.'''   )
														snake_case_		     =							self._write(
														    file_obj=self.path_or_buf			,							orient=lowercase_			,							lines=lowercase_			,							index=lowercase_			,							**self.to_json_kwargs   )
									return written


				def     A_					(						self       :   List[str]			,							lowercase_       :   Dict   ):
									snake_case_     ,snake_case_     ,snake_case_     ,snake_case_     ,snake_case_		     =							args

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


				def     A_					(						self       :   Optional[int]			,							lowercase_       :   BinaryIO			,							lowercase_       :   List[Any]			,							lowercase_       :   str			,							lowercase_       :   int			,							**lowercase_       :   Optional[int]			,							):
									snake_case_		     =							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'''			,							):
																			snake_case_		     =							self._batch_json((offset, orient, lines, index, to_json_kwargs)   )
																			written += file_obj.write(lowercase_   )
									else:
														snake_case_     ,snake_case_		     =							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			,							lowercase_			,							lowercase_   )]			,							)			,							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(lowercase_   )

									return written

