import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_bart import BartTokenizer


A_      =      logging.get_logger(__name__)


A_      =      {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

# See all BART models at https://huggingface.co/models?filter=bart
A_      =      {
    "vocab_file": {
        "facebook/bart-base": "https://huggingface.co/facebook/bart-base/resolve/main/vocab.json",
        "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/vocab.json",
        "facebook/bart-large-mnli": "https://huggingface.co/facebook/bart-large-mnli/resolve/main/vocab.json",
        "facebook/bart-large-cnn": "https://huggingface.co/facebook/bart-large-cnn/resolve/main/vocab.json",
        "facebook/bart-large-xsum": "https://huggingface.co/facebook/bart-large-xsum/resolve/main/vocab.json",
        "yjernite/bart_eli5": "https://huggingface.co/yjernite/bart_eli5/resolve/main/vocab.json",
    },
    "merges_file": {
        "facebook/bart-base": "https://huggingface.co/facebook/bart-base/resolve/main/merges.txt",
        "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/merges.txt",
        "facebook/bart-large-mnli": "https://huggingface.co/facebook/bart-large-mnli/resolve/main/merges.txt",
        "facebook/bart-large-cnn": "https://huggingface.co/facebook/bart-large-cnn/resolve/main/merges.txt",
        "facebook/bart-large-xsum": "https://huggingface.co/facebook/bart-large-xsum/resolve/main/merges.txt",
        "yjernite/bart_eli5": "https://huggingface.co/yjernite/bart_eli5/resolve/main/merges.txt",
    },
    "tokenizer_file": {
        "facebook/bart-base": "https://huggingface.co/facebook/bart-base/resolve/main/tokenizer.json",
        "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/tokenizer.json",
        "facebook/bart-large-mnli": "https://huggingface.co/facebook/bart-large-mnli/resolve/main/tokenizer.json",
        "facebook/bart-large-cnn": "https://huggingface.co/facebook/bart-large-cnn/resolve/main/tokenizer.json",
        "facebook/bart-large-xsum": "https://huggingface.co/facebook/bart-large-xsum/resolve/main/tokenizer.json",
        "yjernite/bart_eli5": "https://huggingface.co/yjernite/bart_eli5/resolve/main/tokenizer.json",
    },
}

A_      =      {
    "facebook/bart-base": 1024,
    "facebook/bart-large": 1024,
    "facebook/bart-large-mnli": 1024,
    "facebook/bart-large-cnn": 1024,
    "facebook/bart-large-xsum": 1024,
    "yjernite/bart_eli5": 1024,
}






class      __lowercase					(    _A    ):
						lowercase					 =     VOCAB_FILES_NAMES
						lowercase					 =     PRETRAINED_VOCAB_FILES_MAP
						lowercase					 =     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						lowercase					 =     ['input_ids', 'attention_mask']
						lowercase					 =     BartTokenizer

						def __init__(	self				:      List[Any]							,							__lowerCamelCase				:      Dict=None							,							__lowerCamelCase				:      Optional[Any]=None							,							__lowerCamelCase				:      Optional[Any]=None							,							__lowerCamelCase				:      int="replace"							,							__lowerCamelCase				:      List[str]="<s>"							,							__lowerCamelCase				:      Dict="</s>"							,							__lowerCamelCase				:      Optional[Any]="</s>"							,							__lowerCamelCase				:      List[str]="<s>"							,							__lowerCamelCase				:      Optional[int]="<unk>"							,							__lowerCamelCase				:      Optional[Any]="<pad>"							,							__lowerCamelCase				:      List[str]="<mask>"							,							__lowerCamelCase				:      Union[str, Any]=False							,							__lowerCamelCase				:      int=True							,							**__lowerCamelCase				:      str							,							)							->		Optional[Any]:


									'''simple docstring'''

									super().__init__(
									    __lowerCamelCase							,							__lowerCamelCase							,							tokenizer_file=__lowerCamelCase							,							errors=__lowerCamelCase							,							bos_token=__lowerCamelCase							,							eos_token=__lowerCamelCase							,							sep_token=__lowerCamelCase							,							cls_token=__lowerCamelCase							,							unk_token=__lowerCamelCase							,							pad_token=__lowerCamelCase							,							mask_token=__lowerCamelCase							,							add_prefix_space=__lowerCamelCase							,							trim_offsets=__lowerCamelCase							,							**__lowerCamelCase							,							)

									lowercase							=		json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()   )
									if pre_tok_state.get('''add_prefix_space'''							,							__lowerCamelCase   ) != add_prefix_space:
												lowercase							=		getattr(__lowerCamelCase							,							pre_tok_state.pop('''type'''   )   )
												lowercase							=		add_prefix_space
												lowercase							=		pre_tok_class(**__lowerCamelCase   )

									lowercase							=		add_prefix_space

									# the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
									lowercase							=		'''post_processor'''
									lowercase							=		getattr(self.backend_tokenizer							,							__lowerCamelCase							,							__lowerCamelCase   )
									if tokenizer_component_instance:
												lowercase							=		json.loads(tokenizer_component_instance.__getstate__()   )

												# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
												if "sep" in state:
															lowercase							=		tuple(state['''sep''']   )
												if "cls" in state:
															lowercase							=		tuple(state['''cls''']   )

												lowercase							=		False

												if state.get('''add_prefix_space'''							,							__lowerCamelCase   ) != add_prefix_space:
															lowercase							=		add_prefix_space
															lowercase							=		True

												if state.get('''trim_offsets'''							,							__lowerCamelCase   ) != trim_offsets:
															lowercase							=		trim_offsets
															lowercase							=		True

												if changes_to_apply:
															lowercase							=		getattr(__lowerCamelCase							,							state.pop('''type'''   )   )
															lowercase							=		component_class(**__lowerCamelCase   )
															setattr(self.backend_tokenizer							,							__lowerCamelCase							,							__lowerCamelCase   )

						@property
						def        __a		(	self				:      List[str]   )							->		str:


									'''simple docstring'''

									if self._mask_token is None:
												if self.verbose:
															logger.error('''Using mask_token, but it is not set yet.'''   )
												return None
									return str(self._mask_token   )

						@mask_token.setter
						def        __a		(	self				:      Optional[int]							,							__lowerCamelCase				:      Any   )							->		int:


									'''simple docstring'''

									lowercase							=		AddedToken(__lowerCamelCase							,							lstrip=__lowerCamelCase							,							rstrip=__lowerCamelCase   ) if isinstance(__lowerCamelCase							,							__lowerCamelCase   ) else value
									lowercase							=		value

						def        __a		(	self				:      Any							,							*__lowerCamelCase				:      Dict							,							**__lowerCamelCase				:      Optional[int]   )							->		BatchEncoding:


									'''simple docstring'''

									lowercase							=		kwargs.get('''is_split_into_words'''							,							__lowerCamelCase   )

									if is_split_into_words and not self.add_prefix_space:
												raise ValueError(
												    f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
												    '''to use it with pretokenized inputs.'''   )

									return super()._batch_encode_plus(*__lowerCamelCase							,							**__lowerCamelCase   )

						def        __a		(	self				:      Any							,							*__lowerCamelCase				:      Any							,							**__lowerCamelCase				:      Optional[int]   )							->		BatchEncoding:


									'''simple docstring'''

									lowercase							=		kwargs.get('''is_split_into_words'''							,							__lowerCamelCase   )

									if is_split_into_words and not self.add_prefix_space:
												raise ValueError(
												    f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
												    '''to use it with pretokenized inputs.'''   )

									return super()._encode_plus(*__lowerCamelCase							,							**__lowerCamelCase   )

						def        __a		(	self				:      str							,							__lowerCamelCase				:      str							,							__lowerCamelCase				:      Optional[str] = None   )							->		Tuple[str]:


									'''simple docstring'''

									lowercase							=		self._tokenizer.model.save(__lowerCamelCase							,							name=__lowerCamelCase   )
									return tuple(__lowerCamelCase   )

						def        __a		(	self				:      Tuple							,							__lowerCamelCase				:      Optional[Any]							,							__lowerCamelCase				:      List[Any]=None   )							->		Union[str, Any]:


									'''simple docstring'''

									lowercase							=		[self.bos_token_id] + token_ids_a + [self.eos_token_id]
									if token_ids_a is None:
												return output

									return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]



						def        __a		(	self				:      int							,							__lowerCamelCase				:      List[int]							,							__lowerCamelCase				:      Optional[List[int]] = None   )							->		List[int]:


									'''simple docstring'''

									lowercase							=		[self.sep_token_id]
									lowercase							=		[self.cls_token_id]

									if token_ids_a is None:
												return len(cls + token_ids_a + sep   ) * [0]
									return len(cls + token_ids_a + sep + sep + token_ids_a + sep   ) * [0]






import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam






class      __lowercase					(    datasets.BeamBasedBuilder    ):

						def        __a		(	self				:      Tuple   )							->		Tuple:


									'''simple docstring'''

									return datasets.DatasetInfo(
									    features=datasets.Features({'''content''': datasets.Value('''string'''   )}   )							,							supervised_keys=__lowerCamelCase							,							)

						def        __a		(	self				:      Union[str, Any]							,							__lowerCamelCase				:      Union[str, Any]							,							__lowerCamelCase				:      Tuple   )							->		int:


									'''simple docstring'''

									return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,							gen_kwargs={'''examples''': get_test_dummy_examples()}   )]



						def        __a		(	self				:      List[Any]							,							__lowerCamelCase				:      Any							,							__lowerCamelCase				:      Union[str, Any]   )							->		Optional[int]:


									'''simple docstring'''

									import apache_beam as beam

									return pipeline | "Load Examples" >> beam.Create(__lowerCamelCase   )






class      __lowercase					(    datasets.BeamBasedBuilder    ):

						def        __a		(	self				:      int   )							->		Any:


									'''simple docstring'''

									return datasets.DatasetInfo(
									    features=datasets.Features({'''a''': datasets.Sequence({'''b''': datasets.Value('''string'''   )}   )}   )							,							supervised_keys=__lowerCamelCase							,							)

						def        __a		(	self				:      Optional[int]							,							__lowerCamelCase				:      List[Any]							,							__lowerCamelCase				:      Optional[int]   )							->		int:


									'''simple docstring'''

									return [
									    datasets.SplitGenerator(name=datasets.Split.TRAIN							,							gen_kwargs={'''examples''': get_test_nested_examples()}   )
									]



						def        __a		(	self				:      Union[str, Any]							,							__lowerCamelCase				:      Tuple							,							__lowerCamelCase				:      Union[str, Any]   )							->		Tuple:


									'''simple docstring'''

									import apache_beam as beam

									return pipeline | "Load Examples" >> beam.Create(__lowerCamelCase   )






def        __UpperCAmelCase    (       )->       Dict:



			"""simple docstring"""







			return [(i, {"content": content}) for i, content in enumerate(['''foo''', '''bar''', '''foobar''']       )]






def        __UpperCAmelCase    (       )->       List[str]:



			"""simple docstring"""







			return [(i, {"a": {"b": [content]}}) for i, content in enumerate(['''foo''', '''bar''', '''foobar''']       )]






class      __lowercase					(    _A    ):

						@require_beam
						def        __a		(	self				:      Optional[int]   )							->		Optional[Any]:


									'''simple docstring'''

									lowercase							=		len(get_test_dummy_examples()   )
									with tempfile.TemporaryDirectory() as tmp_cache_dir:
												lowercase							=		DummyBeamDataset(cache_dir=__lowerCamelCase							,							beam_runner='''DirectRunner'''   )
												builder.download_and_prepare()
												self.assertTrue(
												    os.path.exists(
												        os.path.join(__lowerCamelCase							,							builder.name							,							'''default'''							,							'''0.0.0'''							,							f'{builder.name}-train.arrow'   )   )   )
												self.assertDictEqual(builder.info.features							,							datasets.Features({'''content''': datasets.Value('''string'''   )}   )   )
												lowercase							=		builder.as_dataset()
												self.assertEqual(dset['''train'''].num_rows							,							__lowerCamelCase   )
												self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples							,							__lowerCamelCase   )
												self.assertDictEqual(dset['''train'''][0]							,							get_test_dummy_examples()[0][1]   )
												self.assertDictEqual(
												    dset['''train'''][expected_num_examples - 1]							,							get_test_dummy_examples()[expected_num_examples - 1][1]   )
												self.assertTrue(
												    os.path.exists(os.path.join(__lowerCamelCase							,							builder.name							,							'''default'''							,							'''0.0.0'''							,							'''dataset_info.json'''   )   )   )
												del dset

						@require_beam
						def        __a		(	self				:      Union[str, Any]   )							->		List[Any]:


									'''simple docstring'''

									import apache_beam as beam

									lowercase							=		beam.io.parquetio.WriteToParquet

									lowercase							=		len(get_test_dummy_examples()   )
									with tempfile.TemporaryDirectory() as tmp_cache_dir:
												lowercase							=		DummyBeamDataset(cache_dir=__lowerCamelCase							,							beam_runner='''DirectRunner'''   )
												with patch('''apache_beam.io.parquetio.WriteToParquet'''   ) as write_parquet_mock:
															lowercase							=		partial(__lowerCamelCase							,							num_shards=2   )
															builder.download_and_prepare()
												self.assertTrue(
												    os.path.exists(
												        os.path.join(
												            __lowerCamelCase							,							builder.name							,							'''default'''							,							'''0.0.0'''							,							f'{builder.name}-train-00000-of-00002.arrow'   )   )   )
												self.assertTrue(
												    os.path.exists(
												        os.path.join(
												            __lowerCamelCase							,							builder.name							,							'''default'''							,							'''0.0.0'''							,							f'{builder.name}-train-00000-of-00002.arrow'   )   )   )
												self.assertDictEqual(builder.info.features							,							datasets.Features({'''content''': datasets.Value('''string'''   )}   )   )
												lowercase							=		builder.as_dataset()
												self.assertEqual(dset['''train'''].num_rows							,							__lowerCamelCase   )
												self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples							,							__lowerCamelCase   )
												# Order is not preserved when sharding, so we just check that all the elements are there
												self.assertListEqual(sorted(dset['''train''']['''content''']   )							,							sorted(['''foo''', '''bar''', '''foobar''']   )   )
												self.assertTrue(
												    os.path.exists(os.path.join(__lowerCamelCase							,							builder.name							,							'''default'''							,							'''0.0.0'''							,							'''dataset_info.json'''   )   )   )
												del dset

						@require_beam
						def        __a		(	self				:      int   )							->		int:


									'''simple docstring'''

									with tempfile.TemporaryDirectory() as tmp_cache_dir:
												lowercase							=		DummyBeamDataset(cache_dir=__lowerCamelCase   )
												self.assertRaises(datasets.builder.MissingBeamOptions							,							builder.download_and_prepare   )



						@require_beam
						def        __a		(	self				:      str   )							->		List[Any]:


									'''simple docstring'''

									lowercase							=		len(get_test_nested_examples()   )
									with tempfile.TemporaryDirectory() as tmp_cache_dir:
												lowercase							=		NestedBeamDataset(cache_dir=__lowerCamelCase							,							beam_runner='''DirectRunner'''   )
												builder.download_and_prepare()
												self.assertTrue(
												    os.path.exists(
												        os.path.join(__lowerCamelCase							,							builder.name							,							'''default'''							,							'''0.0.0'''							,							f'{builder.name}-train.arrow'   )   )   )
												self.assertDictEqual(
												    builder.info.features							,							datasets.Features({'''a''': datasets.Sequence({'''b''': datasets.Value('''string'''   )}   )}   )   )
												lowercase							=		builder.as_dataset()
												self.assertEqual(dset['''train'''].num_rows							,							__lowerCamelCase   )
												self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples							,							__lowerCamelCase   )
												self.assertDictEqual(dset['''train'''][0]							,							get_test_nested_examples()[0][1]   )
												self.assertDictEqual(
												    dset['''train'''][expected_num_examples - 1]							,							get_test_nested_examples()[expected_num_examples - 1][1]   )
												self.assertTrue(
												    os.path.exists(os.path.join(__lowerCamelCase							,							builder.name							,							'''default'''							,							'''0.0.0'''							,							'''dataset_info.json'''   )   )   )
												del dset





