






'''simple docstring'''
import unittest

from transformers import PegasusTokenizer, PegasusTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, require_torch, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


UpperCamelCase__ =			get_tests_dir('fixtures/test_sentencepiece_no_bos.model')


@require_sentencepiece
@require_tokenizers
class        _UpperCAmelCase						(		snake_case				,					unittest.TestCase			):
    __lowerCamelCase:   List[str]		       =    PegasusTokenizer
    __lowerCamelCase:   Tuple		       =    PegasusTokenizerFast
    __lowerCamelCase:   List[str]		       =    True
    __lowerCamelCase:   Union[str, Any]		       =    True




    def 				lowerCAmelCase__				(			self  :       Dict			):






       '''simple docstring'''






       super().setUp()

       # We have a SentencePiece fixture for testing
       lowercase_	:    Tuple           =     PegasusTokenizer(a			)
       tokenizer.save_pretrained(self.tmpdirname			)




    @cached_property
    def 				lowerCAmelCase__				(			self  :       Dict			):






       '''simple docstring'''






       return PegasusTokenizer.from_pretrained("google/pegasus-large"			)




    def 				lowerCAmelCase__				(			self  :       Tuple  ,     **a  :       int			):






       '''simple docstring'''






       return PegasusTokenizer.from_pretrained(self.tmpdirname  ,     **a			)




    def 				lowerCAmelCase__				(			self  :       Optional[Any]  ,     a  :       List[str]			):






       '''simple docstring'''






       return ("This is a test", "This is a test")




    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     "</s>"
       lowercase_	:    Tuple           =     1

       self.assertEqual(self.get_tokenizer()._convert_token_to_id(a			)  ,     a			)
       self.assertEqual(self.get_tokenizer()._convert_id_to_token(a			)  ,     a			)




    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       lowercase_	:    str           =     list(self.get_tokenizer().get_vocab().keys()			)

       self.assertEqual(vocab_keys[0]  ,     "<pad>"			)
       self.assertEqual(vocab_keys[1]  ,     "</s>"			)
       self.assertEqual(vocab_keys[-1]  ,     "v"			)
       self.assertEqual(len(a			)  ,     1_1_0_3			)




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       self.assertEqual(self.get_tokenizer().vocab_size  ,     1_1_0_3			)




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     self.rust_tokenizer_class.from_pretrained(self.tmpdirname			)
       lowercase_	:    Optional[int]           =     self.tokenizer_class.from_pretrained(self.tmpdirname			)
       lowercase_	:    str           =     (
           "Let's see which <unk> is the better <unk_token_11> one <mask_1> It seems like this <mask_2> was important"
           " </s> <pad> <pad> <pad>"
       )
       lowercase_	:    str           =     rust_tokenizer([raw_input_str]  ,     return_tensors=a  ,     add_special_tokens=a			).input_ids[0]
       lowercase_	:    Optional[int]           =     py_tokenizer([raw_input_str]  ,     return_tensors=a  ,     add_special_tokens=a			).input_ids[0]
       self.assertListEqual(a  ,     a			)




    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     self._large_tokenizer
       # <mask_1> masks whole sentence while <mask_2> masks single word
       lowercase_	:    Optional[int]           =     "<mask_1> To ensure a <mask_2> flow of bank resolutions."
       lowercase_	:    Any           =     [2, 4_1_3, 6_1_5, 1_1_4, 3, 1_9_7_1, 1_1_3, 1_6_7_9, 1_0_7_1_0, 1_0_7, 1]
       lowercase_	:    List[Any]           =     tokenizer([raw_input_str]  ,     return_tensors=a			).input_ids[0]
       self.assertListEqual(a  ,     a			)




    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       lowercase_	:    str           =     self._large_tokenizer
       # The tracebacks for the following asserts are **better** without messages or self.assertEqual
       assert tokenizer.vocab_size == 9_6_1_0_3
       assert tokenizer.pad_token_id == 0
       assert tokenizer.eos_token_id == 1
       assert tokenizer.offset == 1_0_3
       assert tokenizer.unk_token_id == tokenizer.offset + 2 == 1_0_5
       assert tokenizer.unk_token == "<unk>"
       assert tokenizer.model_max_length == 1_0_2_4
       lowercase_	:    Optional[int]           =     "To ensure a smooth flow of bank resolutions."
       lowercase_	:    Tuple           =     [4_1_3, 6_1_5, 1_1_4, 2_2_9_1, 1_9_7_1, 1_1_3, 1_6_7_9, 1_0_7_1_0, 1_0_7, 1]
       lowercase_	:    str           =     tokenizer([raw_input_str]  ,     return_tensors=a			).input_ids[0]
       self.assertListEqual(a  ,     a			)
       assert tokenizer.convert_ids_to_tokens([0, 1, 2, 3]			) == ["<pad>", "</s>", "<mask_1>", "<mask_2>"]




    @require_torch
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       lowercase_	:    str           =     ["This is going to be way too long." * 1_5_0, "short example"]
       lowercase_	:    Optional[int]           =     ["not super long but more than 5 tokens", "tiny"]
       lowercase_	:    Any           =     self._large_tokenizer(a  ,     padding=a  ,     truncation=a  ,     return_tensors="pt"			)
       lowercase_	:    Union[str, Any]           =     self._large_tokenizer(
           text_target=a  ,     max_length=5  ,     padding=a  ,     truncation=a  ,     return_tensors="pt"			)

       assert batch.input_ids.shape == (2, 1_0_2_4)
       assert batch.attention_mask.shape == (2, 1_0_2_4)
       assert targets["input_ids"].shape == (2, 5)
       assert len(a			) == 2  # input_ids, attention_mask.





    @slow
    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       lowercase_	:    Any           =     {"input_ids": [[3_8_9_7_9, 1_4_3, 1_8_4_8_5, 6_0_6, 1_3_0, 2_6_6_6_9, 8_7_6_8_6, 1_2_1, 5_4_1_8_9, 1_1_2_9, 1_1_1, 2_6_6_6_9, 8_7_6_8_6, 1_2_1, 9_1_1_4, 1_4_7_8_7, 1_2_1, 1_3_2_4_9, 1_5_8, 5_9_2, 9_5_6, 1_2_1, 1_4_6_2_1, 3_1_5_7_6, 1_4_3, 6_2_6_1_3, 1_0_8, 9_6_8_8, 9_3_0, 4_3_4_3_0, 1_1_5_6_2, 6_2_6_1_3, 3_0_4, 1_0_8, 1_1_4_4_3, 8_9_7, 1_0_8, 9_3_1_4, 1_7_4_1_5, 6_3_3_9_9, 1_0_8, 1_1_4_4_3, 7_6_1_4, 1_8_3_1_6, 1_1_8, 4_2_8_4, 7_1_4_8, 1_2_4_3_0, 1_4_3, 1_4_0_0, 2_5_7_0_3, 1_5_8, 1_1_1, 4_2_8_4, 7_1_4_8, 1_1_7_7_2, 1_4_3, 2_1_2_9_7, 1_0_6_4, 1_5_8, 1_2_2, 2_0_4, 3_5_0_6, 1_7_5_4, 1_1_3_3, 1_4_7_8_7, 1_5_8_1, 1_1_5, 3_3_2_2_4, 4_4_8_2, 1_1_1, 1_3_5_5, 1_1_0, 2_9_1_7_3, 3_1_7, 5_0_8_3_3, 1_0_8, 2_0_1_4_7, 9_4_6_6_5, 1_1_1, 7_7_1_9_8, 1_0_7, 1], [1_1_0, 6_2_6_1_3, 1_1_7, 6_3_8, 1_1_2, 1_1_3_3, 1_2_1, 2_0_0_9_8, 1_3_5_5, 7_9_0_5_0, 1_3_8_7_2, 1_3_5, 1_5_9_6, 5_3_5_4_1, 1_3_5_2, 1_4_1, 1_3_0_3_9, 5_5_4_2, 1_2_4, 3_0_2, 5_1_8, 1_1_1, 2_6_8, 2_9_5_6, 1_1_5, 1_4_9, 4_4_2_7, 1_0_7, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1_3_9, 1_2_3_5, 2_7_9_9, 1_8_2_8_9, 1_7_7_8_0, 2_0_4, 1_0_9, 9_4_7_4, 1_2_9_6, 1_0_7, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], "attention_mask": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
       # fmt: on

       self.tokenizer_integration_test_util(
           expected_encoding=a  ,     model_name="google/bigbird-pegasus-large-arxiv"  ,     revision="ba85d0851d708441f91440d509690f1ab6353415"  ,     )


@require_sentencepiece
@require_tokenizers
class        _UpperCAmelCase						(		snake_case				,					unittest.TestCase			):
    __lowerCamelCase:   int		       =    PegasusTokenizer
    __lowerCamelCase:   List[str]		       =    PegasusTokenizerFast
    __lowerCamelCase:   Any		       =    True
    __lowerCamelCase:   Optional[int]		       =    True




    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       super().setUp()

       # We have a SentencePiece fixture for testing
       lowercase_	:    Any           =     PegasusTokenizer(a  ,     offset=0  ,     mask_token_sent=a  ,     mask_token="[MASK]"			)
       tokenizer.save_pretrained(self.tmpdirname			)




    @cached_property
    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       return PegasusTokenizer.from_pretrained("google/bigbird-pegasus-large-arxiv"			)




    def 				lowerCAmelCase__				(			self  :       str  ,     **a  :       Any			):






       '''simple docstring'''






       return PegasusTokenizer.from_pretrained(self.tmpdirname  ,     **a			)




    def 				lowerCAmelCase__				(			self  :       Optional[int]  ,     a  :       Tuple			):






       '''simple docstring'''






       return ("This is a test", "This is a test")




    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       lowercase_	:    Any           =     self.rust_tokenizer_class.from_pretrained(self.tmpdirname			)
       lowercase_	:    Dict           =     self.tokenizer_class.from_pretrained(self.tmpdirname			)
       lowercase_	:    Any           =     (
           "Let's see which <unk> is the better <unk_token> one [MASK] It seems like this [MASK] was important </s>"
           " <pad> <pad> <pad>"
       )
       lowercase_	:    Dict           =     rust_tokenizer([raw_input_str]  ,     return_tensors=a  ,     add_special_tokens=a			).input_ids[0]
       lowercase_	:    Optional[int]           =     py_tokenizer([raw_input_str]  ,     return_tensors=a  ,     add_special_tokens=a			).input_ids[0]
       self.assertListEqual(a  ,     a			)




    @require_torch
    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    Optional[int]           =     ["This is going to be way too long." * 1_0_0_0, "short example"]
       lowercase_	:    List[str]           =     ["not super long but more than 5 tokens", "tiny"]
       lowercase_	:    str           =     self._large_tokenizer(a  ,     padding=a  ,     truncation=a  ,     return_tensors="pt"			)
       lowercase_	:    Union[str, Any]           =     self._large_tokenizer(
           text_target=a  ,     max_length=5  ,     padding=a  ,     truncation=a  ,     return_tensors="pt"			)

       assert batch.input_ids.shape == (2, 4_0_9_6)
       assert batch.attention_mask.shape == (2, 4_0_9_6)
       assert targets["input_ids"].shape == (2, 5)
       assert len(a			) == 2  # input_ids, attention_mask.





    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    Any           =     (
           "This is an example string that is used to test the original TF implementation against the HF"
           " implementation"
       )

       lowercase_	:    Dict           =     self._large_tokenizer(a			).input_ids

       self.assertListEqual(
           a  ,     [1_8_2, 1_1_7, 1_4_2, 5_8_7, 4_2_1_1, 1_2_0, 1_1_7, 2_6_3, 1_1_2, 8_0_4, 1_0_9, 8_5_6, 2_5_0_1_6, 3_1_3_7, 4_6_4, 1_0_9, 2_6_9_5_5, 3_1_3_7, 1]  ,     )












'''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        _UpperCAmelCase						(		snake_case			):




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






       '''simple docstring'''






       super().__init__(
           a  ,     split=a  ,     features=a  ,     cache_dir=a  ,     keep_in_memory=a  ,     streaming=a  ,     num_proc=a  ,     **a  ,     )
       lowercase_	:    str           =     field
       lowercase_	:    Optional[Any]           =     path_or_paths if isinstance(a  ,     a			) else {self.split: path_or_paths}
       lowercase_	:    Any           =     Json(
           cache_dir=a  ,     data_files=a  ,     features=a  ,     field=a  ,     **a  ,     )





    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       if self.streaming:
          lowercase_	:    Any           =     self.builder.as_streaming_dataset(split=self.split			)
       # Build regular (map-style) dataset
       else:
          lowercase_	:    Dict           =     None
          lowercase_	:    Optional[int]           =     None
          lowercase_	:    str           =     None
          lowercase_	:    str           =     None

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


class        _UpperCAmelCase						:




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






       '''simple docstring'''






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

       lowercase_	:    Dict           =     dataset
       lowercase_	:    Optional[int]           =     path_or_buf
       lowercase_	:    List[Any]           =     batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
       lowercase_	:    Optional[Any]           =     num_proc
       lowercase_	:    List[Any]           =     "utf-8"
       lowercase_	:    List[str]           =     to_json_kwargs




    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    str           =     self.to_json_kwargs.pop("path_or_buf"  ,     a			)
       lowercase_	:    Any           =     self.to_json_kwargs.pop("orient"  ,     "records"			)
       lowercase_	:    Any           =     self.to_json_kwargs.pop("lines"  ,     True if orient == "records" else False			)
       lowercase_	:    List[str]           =     self.to_json_kwargs.pop("index"  ,     False if orient in ["split", "table"] else True			)
       lowercase_	:    int           =     self.to_json_kwargs.pop("compression"  ,     a			)

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




    def 				lowerCAmelCase__				(			self  :       Optional[int]  ,     a  :       List[str]			):






       '''simple docstring'''






       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    List[Any]           =     args

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





    def 				lowerCAmelCase__				(			self  :       int  ,     a  :       BinaryIO  ,     a  :       int  ,     a  :       str  ,     a  :       Union[str, Any]  ,     **a  :       str  ,     ):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     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"  ,     ):
             lowercase_	:    Dict           =     self._batch_json((offset, orient, lines, index, to_json_kwargs)			)
             written += file_obj.write(a			)
       else:
          lowercase_    ,		lowercase_	:    Any           =     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  ,     a  ,     a			)]  ,     )  ,     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(a			)

       return written




