






import os
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


SCREAMING_SNAKE_CASE__   :      int				=		logging.get_logger(__name__)
SCREAMING_SNAKE_CASE__   :      Any				=		{"""vocab_file""": """spiece.model"""}

SCREAMING_SNAKE_CASE__   :      str				=		{
    """vocab_file""": {
        """albert-base-v1""": """https://huggingface.co/albert-base-v1/resolve/main/spiece.model""",
        """albert-large-v1""": """https://huggingface.co/albert-large-v1/resolve/main/spiece.model""",
        """albert-xlarge-v1""": """https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model""",
        """albert-xxlarge-v1""": """https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model""",
        """albert-base-v2""": """https://huggingface.co/albert-base-v2/resolve/main/spiece.model""",
        """albert-large-v2""": """https://huggingface.co/albert-large-v2/resolve/main/spiece.model""",
        """albert-xlarge-v2""": """https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model""",
        """albert-xxlarge-v2""": """https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model""",
    }
}

SCREAMING_SNAKE_CASE__   :      Union[str, Any]				=		{
    """albert-base-v1""": 5_12,
    """albert-large-v1""": 5_12,
    """albert-xlarge-v1""": 5_12,
    """albert-xxlarge-v1""": 5_12,
    """albert-base-v2""": 5_12,
    """albert-large-v2""": 5_12,
    """albert-xlarge-v2""": 5_12,
    """albert-xxlarge-v2""": 5_12,
}

SCREAMING_SNAKE_CASE__   :      List[str]				=		"""▁"""
class     lowerCamelCase_					(   lowerCamelCase			):
				a__       							=	VOCAB_FILES_NAMES
				a__       							=	PRETRAINED_VOCAB_FILES_MAP
				a__       							=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES







				def __init__(    self							,       __lowerCAmelCase							,       __lowerCAmelCase=True							,       __lowerCAmelCase=True							,       __lowerCAmelCase=False							,       __lowerCAmelCase="[CLS]"							,       __lowerCAmelCase="[SEP]"							,       __lowerCAmelCase="<unk>"							,       __lowerCAmelCase="[SEP]"							,       __lowerCAmelCase="<pad>"							,       __lowerCAmelCase="[CLS]"							,       __lowerCAmelCase="[MASK]"							,       __lowerCAmelCase = None							,       **__lowerCAmelCase							,       ):
								"""simple docstring"""
								# Mask token behave like a normal word, i.e. include the space before it and
								# is included in the raw text, there should be a match in a non-normalized sentence.
								__magic_name__		:Any   =   (
								    AddedToken(__lowerCAmelCase							,       lstrip=__lowerCAmelCase							,       rstrip=__lowerCAmelCase							,       normalized=__lowerCAmelCase  )
								    if isinstance(__lowerCAmelCase							,       __lowerCAmelCase  )
								    else mask_token
								)

								__magic_name__		:int   =   {} if sp_model_kwargs is None else sp_model_kwargs

								super().__init__(
								    do_lower_case=__lowerCAmelCase							,       remove_space=__lowerCAmelCase							,       keep_accents=__lowerCAmelCase							,       bos_token=__lowerCAmelCase							,       eos_token=__lowerCAmelCase							,       unk_token=__lowerCAmelCase							,       sep_token=__lowerCAmelCase							,       pad_token=__lowerCAmelCase							,       cls_token=__lowerCAmelCase							,       mask_token=__lowerCAmelCase							,       sp_model_kwargs=self.sp_model_kwargs							,       **__lowerCAmelCase							,       )

								__magic_name__		:Optional[Any]   =   do_lower_case
								__magic_name__		:Any   =   remove_space
								__magic_name__		:List[str]   =   keep_accents
								__magic_name__		:int   =   vocab_file

								__magic_name__		:Optional[int]   =   spm.SentencePieceProcessor(**self.sp_model_kwargs  )
								self.sp_model.Load(__lowerCAmelCase  )







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return len(self.sp_model  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   {self.convert_ids_to_tokens(__lowerCAmelCase  ): i for i in range(self.vocab_size  )}
								vocab.update(self.added_tokens_encoder  )
								return vocab







				def __getstate__(    self  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   self.__dict__.copy()
								__magic_name__		:Union[str, Any]   =   None
								return state







				def __setstate__(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   d

								# for backward compatibility
								if not hasattr(self							,       '''sp_model_kwargs'''  ):
												__magic_name__		:Optional[Any]   =   {}

								__magic_name__		:int   =   spm.SentencePieceProcessor(**self.sp_model_kwargs  )
								self.sp_model.Load(self.vocab_file  )







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								if self.remove_space:
												__magic_name__		:List[str]   =   ''' '''.join(inputs.strip().split()  )
								else:
												__magic_name__		:Optional[int]   =   inputs
								__magic_name__		:Dict   =   outputs.replace('''``'''							,       '''"'''  ).replace('''\'\''''							,       '''"'''  )

								if not self.keep_accents:
												__magic_name__		:str   =   unicodedata.normalize('''NFKD'''							,       __lowerCAmelCase  )
												__magic_name__		:Any   =   ''''''.join([c for c in outputs if not unicodedata.combining(__lowerCAmelCase  )]  )
								if self.do_lower_case:
												__magic_name__		:Optional[Any]   =   outputs.lower()

								return outputs







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:Optional[int]   =   self.preprocess_text(__lowerCAmelCase  )
								__magic_name__		:Dict   =   self.sp_model.encode(__lowerCAmelCase							,       out_type=__lowerCAmelCase  )
								__magic_name__		:Any   =   []
								for piece in pieces:
												if len(__lowerCAmelCase  ) > 1 and piece[-1] == str(''','''  ) and piece[-2].isdigit():
																__magic_name__		:List[str]   =   self.sp_model.EncodeAsPieces(piece[:-1].replace(__lowerCAmelCase							,       ''''''  )  )
																if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
																				if len(cur_pieces[0]  ) == 1:
																								__magic_name__		:str   =   cur_pieces[1:]
																				else:
																								__magic_name__		:List[Any]   =   cur_pieces[0][1:]
																cur_pieces.append(piece[-1]  )
																new_pieces.extend(__lowerCAmelCase  )
												else:
																new_pieces.append(__lowerCAmelCase  )

								return new_pieces







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								return self.sp_model.PieceToId(__lowerCAmelCase  )







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								return self.sp_model.IdToPiece(__lowerCAmelCase  )







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   []
								__magic_name__		:List[str]   =   ''''''
								__magic_name__		:List[Any]   =   False
								for token in tokens:
												# make sure that special tokens are not decoded using sentencepiece model
												if token in self.all_special_tokens:
																if not prev_is_special:
																				out_string += " "
																out_string += self.sp_model.decode(__lowerCAmelCase  ) + token
																__magic_name__		:Union[str, Any]   =   True
																__magic_name__		:Tuple   =   []
												else:
																current_sub_tokens.append(__lowerCAmelCase  )
																__magic_name__		:Optional[int]   =   False
								out_string += self.sp_model.decode(__lowerCAmelCase  )
								return out_string.strip()







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase = None  ):
								"""simple docstring"""
								__magic_name__		:Any   =   [self.sep_token_id]
								__magic_name__		:List[Any]   =   [self.cls_token_id]
								if token_ids_a is None:
												return cls + token_ids_a + sep
								return cls + token_ids_a + sep + token_ids_a + sep







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase = None							,       __lowerCAmelCase = False  ):
								"""simple docstring"""

								if already_has_special_tokens:
												return super().get_special_tokens_mask(
												    token_ids_a=__lowerCAmelCase							,       token_ids_a=__lowerCAmelCase							,       already_has_special_tokens=__lowerCAmelCase  )

								if token_ids_a is not None:
												return [1] + ([0] * len(__lowerCAmelCase  )) + [1] + ([0] * len(__lowerCAmelCase  )) + [1]
								return [1] + ([0] * len(__lowerCAmelCase  )) + [1]







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase = None  ):
								"""simple docstring"""
								__magic_name__		:Union[str, Any]   =   [self.sep_token_id]
								__magic_name__		:str   =   [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  ) * [0] + len(token_ids_a + sep  ) * [1]







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase = None  ):
								"""simple docstring"""
								if not os.path.isdir(__lowerCAmelCase  ):
												logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''  )
												return
								__magic_name__		:int   =   os.path.join(
								    __lowerCAmelCase							,       (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']  )

								if os.path.abspath(self.vocab_file  ) != os.path.abspath(__lowerCAmelCase  ) and os.path.isfile(self.vocab_file  ):
												copyfile(self.vocab_file							,       __lowerCAmelCase  )
								elif not os.path.isfile(self.vocab_file  ):
												with open(__lowerCAmelCase							,       '''wb'''  ) as fi:
																__magic_name__		:int   =   self.sp_model.serialized_model_proto()
																fi.write(__lowerCAmelCase  )

								return (out_vocab_file,)














import pyarrow.parquet as pq
import pytest

from datasets import Audio, Dataset, DatasetDict, Features, NamedSplit, Sequence, Value, config
from datasets.features.image import Image
from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter, get_writer_batch_size

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases

def   __lowercase       (			snake_case,	snake_case			):







				"""simple docstring"""






				assert isinstance(snake_case,	snake_case			)
				assert dataset.num_rows == 4
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_1", "col_2", "col_3"]
				for feature, expected_dtype in expected_features.items():
								assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize('''keep_in_memory''',	[False, True]			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:Tuple   =   tmp_path / '''cache'''
				__magic_name__		:int   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								__magic_name__		:Optional[Any]   =   ParquetDatasetReader(snake_case,	cache_dir=snake_case,	keep_in_memory=snake_case			).read()
				_check_parquet_dataset(snake_case,	snake_case			)

@pytest.mark.parametrize(
    '''features''',	[
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ],	)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:List[str]   =   tmp_path / '''cache'''
				__magic_name__		:int   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:Tuple   =   features.copy() if features else default_expected_features
				__magic_name__		:Union[str, Any]   =   (
				    Features({feature: Value(snake_case			) for feature, dtype in features.items()}			) if features is not None else None
				)
				__magic_name__		:int   =   ParquetDatasetReader(snake_case,	features=snake_case,	cache_dir=snake_case			).read()
				_check_parquet_dataset(snake_case,	snake_case			)

@pytest.mark.parametrize('''split''',	[None, NamedSplit('''train'''			), '''train''', '''test''']			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:str   =   tmp_path / '''cache'''
				__magic_name__		:List[Any]   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:int   =   ParquetDatasetReader(snake_case,	cache_dir=snake_case,	split=snake_case			).read()
				_check_parquet_dataset(snake_case,	snake_case			)
				assert dataset.split == split if split else "train"

@pytest.mark.parametrize('''path_type''',	[str, list]			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				if issubclass(snake_case,	snake_case			):
								__magic_name__		:Union[str, Any]   =   parquet_path
				elif issubclass(snake_case,	snake_case			):
								__magic_name__		:Union[str, Any]   =   [parquet_path]
				__magic_name__		:Optional[int]   =   tmp_path / '''cache'''
				__magic_name__		:Optional[int]   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:str   =   ParquetDatasetReader(snake_case,	cache_dir=snake_case			).read()
				_check_parquet_dataset(snake_case,	snake_case			)

def   __lowercase       (			snake_case,	snake_case,	snake_case=("train",)			):







				"""simple docstring"""






				assert isinstance(snake_case,	snake_case			)
				for split in splits:
								__magic_name__		:Optional[Any]   =   dataset_dict[split]
								assert dataset.num_rows == 4
								assert dataset.num_columns == 3
								assert dataset.column_names == ["col_1", "col_2", "col_3"]
								for feature, expected_dtype in expected_features.items():
												assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize('''keep_in_memory''',	[False, True]			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:Any   =   tmp_path / '''cache'''
				__magic_name__		:Optional[int]   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								__magic_name__		:Tuple   =   ParquetDatasetReader(
								    {'''train''': parquet_path},	cache_dir=snake_case,	keep_in_memory=snake_case			).read()
				_check_parquet_datasetdict(snake_case,	snake_case			)

@pytest.mark.parametrize(
    '''features''',	[
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ],	)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:Optional[Any]   =   tmp_path / '''cache'''
				__magic_name__		:Dict   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:int   =   features.copy() if features else default_expected_features
				__magic_name__		:List[Any]   =   (
				    Features({feature: Value(snake_case			) for feature, dtype in features.items()}			) if features is not None else None
				)
				__magic_name__		:Optional[int]   =   ParquetDatasetReader({'''train''': parquet_path},	features=snake_case,	cache_dir=snake_case			).read()
				_check_parquet_datasetdict(snake_case,	snake_case			)

@pytest.mark.parametrize('''split''',	[None, NamedSplit('''train'''			), '''train''', '''test''']			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				if split:
								__magic_name__		:Dict   =   {split: parquet_path}
				else:
								__magic_name__		:Optional[int]   =   '''train'''
								__magic_name__		:Dict   =   {'''train''': parquet_path, '''test''': parquet_path}
				__magic_name__		:List[Any]   =   tmp_path / '''cache'''
				__magic_name__		:Optional[Any]   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:Optional[Any]   =   ParquetDatasetReader(snake_case,	cache_dir=snake_case			).read()
				_check_parquet_datasetdict(snake_case,	snake_case,	splits=list(path.keys()			)			)
				assert all(dataset[split].split == split for split in path.keys()			)

def   __lowercase       (			snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:str   =   ParquetDatasetWriter(snake_case,	tmp_path / '''foo.parquet'''			)
				assert writer.write() > 0
				__magic_name__		:List[Any]   =   pq.ParquetFile(tmp_path / '''foo.parquet'''			)
				__magic_name__		:List[Any]   =   pf.read()
				assert dataset.data.table == output_table

def   __lowercase       (			snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:List[str]   =   str(shared_datadir / '''test_image_rgb.jpg'''			)
				__magic_name__		:Tuple   =   {'''image''': [image_path]}
				__magic_name__		:List[Any]   =   Features({'''image''': Image()}			)
				__magic_name__		:Tuple   =   Dataset.from_dict(snake_case,	features=snake_case			)
				__magic_name__		:Union[str, Any]   =   ParquetDatasetWriter(snake_case,	tmp_path / '''foo.parquet'''			)
				assert writer.write() > 0

				__magic_name__		:List[str]   =   Dataset.from_parquet(str(tmp_path / '''foo.parquet'''			)			)
				assert dataset.features == reloaded_dataset.features

				__magic_name__		:List[str]   =   ParquetDatasetReader(str(tmp_path / '''foo.parquet'''			),	streaming=snake_case			).read()
				assert dataset.features == reloaded_iterable_dataset.features

@pytest.mark.parametrize(
    '''feature, expected''',	[
        (Features({'''foo''': Value('''int32'''			)}			), None),
        (Features({'''image''': Image(), '''foo''': Value('''int32'''			)}			), config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS),
        (Features({'''nested''': Sequence(Audio()			)}			), config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS),
    ],	)
def   __lowercase       (			snake_case,	snake_case			):







				"""simple docstring"""






				assert get_writer_batch_size(snake_case			) == expected






