


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

import sentencepiece as spm

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


SCREAMING_SNAKE_CASE :   Optional[Any]							      =				logging.get_logger(__name__)

SCREAMING_SNAKE_CASE :   Union[str, Any]							      =				"▁"

SCREAMING_SNAKE_CASE :   Union[str, Any]							      =				{"vocab_file": "sentencepiece.bpe.model"}

SCREAMING_SNAKE_CASE :   List[Any]							      =				{
    "vocab_file": {
        "facebook/mbart-large-50-one-to-many-mmt": (
            "https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model"
        ),
    }
}

SCREAMING_SNAKE_CASE :   str							      =				{
    "facebook/mbart-large-50-one-to-many-mmt": 1024,
}

# fmt: off
SCREAMING_SNAKE_CASE :   List[str]							      =				["ar_AR", "cs_CZ", "de_DE", "en_XX", "es_XX", "et_EE", "fi_FI", "fr_XX", "gu_IN", "hi_IN", "it_IT", "ja_XX", "kk_KZ", "ko_KR", "lt_LT", "lv_LV", "my_MM", "ne_NP", "nl_XX", "ro_RO", "ru_RU", "si_LK", "tr_TR", "vi_VN", "zh_CN", "af_ZA", "az_AZ", "bn_IN", "fa_IR", "he_IL", "hr_HR", "id_ID", "ka_GE", "km_KH", "mk_MK", "ml_IN", "mn_MN", "mr_IN", "pl_PL", "ps_AF", "pt_XX", "sv_SE", "sw_KE", "ta_IN", "te_IN", "th_TH", "tl_XX", "uk_UA", "ur_PK", "xh_ZA", "gl_ES", "sl_SI"]






class _lowerCamelCase(		_a     ):
       lowercase_				:			str									=    VOCAB_FILES_NAMES
       lowercase_				:			Union[str, Any]									=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       lowercase_				:			int									=    PRETRAINED_VOCAB_FILES_MAP
       lowercase_				:			str									=    ["""input_ids""", """attention_mask"""]

       lowercase_				:			List[int]									=    []
       lowercase_				:			List[int]									=    []
       def __init__(      self,   lowerCamelCase,   lowerCamelCase=None,   lowerCamelCase=None,   lowerCamelCase="</s>",   lowerCamelCase="</s>",   lowerCamelCase="<s>",   lowerCamelCase="<unk>",   lowerCamelCase="<pad>",   lowerCamelCase="<mask>",   lowerCamelCase = None,   **lowerCamelCase,   ) ->     None:




           """simple docstring"""

           _lowercase :							List[str]   	=	AddedToken(lowerCamelCase,   lstrip=lowerCamelCase,   rstrip=lowerCamelCase) if isinstance(lowerCamelCase,   lowerCamelCase) else mask_token

           _lowercase :							Any   	=	{} if sp_model_kwargs is None else sp_model_kwargs

           _lowercase :							Optional[Any]   	=	kwargs.get('additional_special_tokens',   [])
           kwargs["additional_special_tokens"] += [
               code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
           ]

           super().__init__(
               src_lang=lowerCamelCase,   tgt_lang=lowerCamelCase,   eos_token=lowerCamelCase,   unk_token=lowerCamelCase,   sep_token=lowerCamelCase,   cls_token=lowerCamelCase,   pad_token=lowerCamelCase,   mask_token=lowerCamelCase,   sp_model_kwargs=self.sp_model_kwargs,   **lowerCamelCase,   )

           _lowercase :							Optional[Any]   	=	spm.SentencePieceProcessor(**self.sp_model_kwargs)
           self.sp_model.Load(str(lowerCamelCase))
           _lowercase :							Any   	=	vocab_file

           # Original fairseq vocab and spm vocab must be "aligned":
           # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
           # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
           # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
           # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

           # Mimic fairseq token-to-id alignment for the first 4 token
           _lowercase :							Dict   	=	{'<s>': 0, '<pad>': 1, '</s>': 2, '<unk>': 3}

           # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
           _lowercase :							List[str]   	=	1

           _lowercase :							int   	=	len(self.sp_model)
           _lowercase :							str   	=	{
               code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(lowerCamelCase)
           }
           _lowercase :							Optional[Any]   	=	{v: k for k, v in self.lang_code_to_id.items()}
           _lowercase :							List[str]   	=	len(self.sp_model) + len(self.lang_code_to_id) + self.fairseq_offset

           self.fairseq_tokens_to_ids.update(self.lang_code_to_id)
           _lowercase :							List[str]   	=	{v: k for k, v in self.fairseq_tokens_to_ids.items()}

           _lowercase :							Union[str, Any]   	=	src_lang if src_lang is not None else 'en_XX'
           _lowercase :							Tuple   	=	self.lang_code_to_id[self._src_lang]
           _lowercase :							int   	=	tgt_lang
           self.set_src_lang_special_tokens(self._src_lang)
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           return len(self.sp_model) + len(self.lang_code_to_id) + self.fairseq_offset + 1  # Plus 1 for the mask token
       @property
       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           return self._src_lang
       @src_lang.setter
       def     UpperCamelCase  (      self,   lowerCamelCase) ->     None:




           """simple docstring"""

           _lowercase :							Dict   	=	new_src_lang
           self.set_src_lang_special_tokens(self._src_lang)
       def __getstate__(      self) ->     Dict:




           """simple docstring"""

           _lowercase :							Tuple   	=	self.__dict__.copy()
           _lowercase :							Any   	=	None
           return state
       def __setstate__(      self,   lowerCamelCase) ->     None:




           """simple docstring"""

           _lowercase :							Optional[int]   	=	d

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

           _lowercase :							Union[str, Any]   	=	spm.SentencePieceProcessor(**self.sp_model_kwargs)
           self.sp_model.Load(self.vocab_file)
       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           _lowercase :							Dict   	=	{self.convert_ids_to_tokens(lowerCamelCase): i for i in range(self.vocab_size)}
           vocab.update(self.added_tokens_encoder)
           return vocab
       def     UpperCamelCase  (      self,   lowerCamelCase) ->     List[str]:




           """simple docstring"""

           return self.sp_model.encode(lowerCamelCase,   out_type=lowerCamelCase)
       def     UpperCamelCase  (      self,   lowerCamelCase) ->     int:




           """simple docstring"""

           if token in self.fairseq_tokens_to_ids:
               return self.fairseq_tokens_to_ids[token]
           _lowercase :							Union[str, Any]   	=	self.sp_model.PieceToId(lowerCamelCase)

           # Need to return unknown token if the SP model returned 0
           return spm_id + self.fairseq_offset if spm_id else self.unk_token_id
       def     UpperCamelCase  (      self,   lowerCamelCase) ->     str:




           """simple docstring"""

           if index in self.fairseq_ids_to_tokens:
               return self.fairseq_ids_to_tokens[index]
           return self.sp_model.IdToPiece(index - self.fairseq_offset)
       def     UpperCamelCase  (      self,   lowerCamelCase) ->     str:




           """simple docstring"""

           _lowercase :							Dict   	=	[]
           _lowercase :							Any   	=	''
           _lowercase :							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
                   _lowercase :							Tuple   	=	True
                   _lowercase :							str   	=	[]
               else:
                   current_sub_tokens.append(lowerCamelCase)
                   _lowercase :							Any   	=	False
           out_string += self.sp_model.decode(lowerCamelCase)
           return out_string.strip()
       def     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase = None) ->     Tuple[str]:




           """simple docstring"""

           if not os.path.isdir(lowerCamelCase):
               logger.error(F'''Vocabulary path ({save_directory}) should be a directory''')
               return
           _lowercase :							List[Any]   	=	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:
                   _lowercase :							Union[str, Any]   	=	self.sp_model.serialized_model_proto()
                   fi.write(lowerCamelCase)

           return (out_vocab_file,)
       def     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase = None,   lowerCamelCase = False) ->     List[int]:




           """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)

           _lowercase :							Optional[Any]   	=	[1] * len(self.prefix_tokens)
           _lowercase :							str   	=	[1] * len(self.suffix_tokens)
           if token_ids_a is None:
               return prefix_ones + ([0] * len(lowerCamelCase)) + suffix_ones
           return prefix_ones + ([0] * len(lowerCamelCase)) + ([0] * len(lowerCamelCase)) + suffix_ones
       def     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase = None) ->     List[int]:




           """simple docstring"""

           if token_ids_a is None:
               return self.prefix_tokens + token_ids_a + self.suffix_tokens
           # We don't expect to process pairs, but leave the pair logic for API consistency
           return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens
       def     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   **lowerCamelCase) ->     List[Any]:




           """simple docstring"""

           if src_lang is None or tgt_lang is None:
               raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model')
           _lowercase :							Tuple   	=	src_lang
           _lowercase :							List[str]   	=	self(lowerCamelCase,   add_special_tokens=lowerCamelCase,   return_tensors=lowerCamelCase,   **lowerCamelCase)
           _lowercase :							Tuple   	=	self.convert_tokens_to_ids(lowerCamelCase)
           _lowercase :							str   	=	tgt_lang_id
           return inputs
       def     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase = "en_XX",   lowerCamelCase = None,   lowerCamelCase = "ro_RO",   **lowerCamelCase,   ) ->     BatchEncoding:




           """simple docstring"""

           _lowercase :							List[Any]   	=	src_lang
           _lowercase :							str   	=	tgt_lang
           return super().prepare_seqaseq_batch(lowerCamelCase,   lowerCamelCase,   **lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           return self.set_src_lang_special_tokens(self.src_lang)
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

           return self.set_tgt_lang_special_tokens(self.tgt_lang)
       def     UpperCamelCase  (      self,   lowerCamelCase) ->     None:




           """simple docstring"""

           _lowercase :							Any   	=	self.lang_code_to_id[src_lang]
           _lowercase :							Optional[int]   	=	[self.cur_lang_code_id]
           _lowercase :							List[Any]   	=	[self.eos_token_id]




       def     UpperCamelCase  (      self,   lowerCamelCase) ->     None:




           """simple docstring"""

           _lowercase :							int   	=	self.lang_code_to_id[tgt_lang]
           _lowercase :							int   	=	[self.cur_lang_code_id]
           _lowercase :							Dict   	=	[self.eos_token_id]






import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict



@pytest.mark.parametrize(
    'files'							,				[
        ['full:README.md', 'dataset_infos.json'],
        ['empty:README.md', 'dataset_infos.json'],
        ['dataset_infos.json'],
        ['full:README.md'],
    ]							,				)
def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_				)					->		List[str]:
    _lowercase :							Any   	=	tmp_path_factory.mktemp('dset_infos_dir'				)
    if "full:README.md" in files:
        with open(dataset_infos_dir / 'README.md'							,				'w'				) as f:
            f.write('---\ndataset_info:\n  dataset_size: 42\n---'				)
    if "empty:README.md" in files:
        with open(dataset_infos_dir / 'README.md'							,				'w'				) as f:
            f.write(''				)
    # we want to support dataset_infos.json for backward compatibility
    if "dataset_infos.json" in files:
        with open(dataset_infos_dir / 'dataset_infos.json'							,				'w'				) as f:
            f.write('{"default": {"dataset_size": 42}}'				)
    _lowercase :							Any   	=	DatasetInfosDict.from_directory(lowerCamelCase_				)
    assert dataset_infos
    assert dataset_infos["default"].dataset_size == 42



@pytest.mark.parametrize(
    'dataset_info'							,				[
        DatasetInfo(),
        DatasetInfo(
            description='foo'							,				features=Features({'a': Value('int32'				)}				)							,				builder_name='builder'							,				config_name='config'							,				version='1.0.0'							,				splits=[{'name': 'train'}]							,				download_size=42							,				),
    ]							,				)
def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_				)					->		Optional[int]:
    _lowercase :							Union[str, Any]   	=	str(lowerCamelCase_				)
    dataset_info.write_to_directory(lowerCamelCase_				)
    _lowercase :							List[str]   	=	DatasetInfo.from_directory(lowerCamelCase_				)
    assert dataset_info == reloaded
    assert os.path.exists(os.path.join(lowerCamelCase_							,				'dataset_info.json'				)				)



def        UpperCamelCase_(   )					->		int:
    _lowercase :							Tuple   	=	DatasetInfo(
        description='foo'							,				citation='bar'							,				homepage='https://foo.bar'							,				license='CC0'							,				features=Features({'a': Value('int32'				)}				)							,				post_processed={}							,				supervised_keys=()							,				task_templates=[]							,				builder_name='builder'							,				config_name='config'							,				version='1.0.0'							,				splits=[{'name': 'train', 'num_examples': 42}]							,				download_checksums={}							,				download_size=1337							,				post_processing_size=442							,				dataset_size=1234							,				size_in_bytes=1337 + 442 + 1234							,				)
    _lowercase :							Optional[int]   	=	dataset_info._to_yaml_dict()
    assert sorted(lowerCamelCase_				) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML				)
    for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
        assert key in dataset_info_yaml_dict
        assert isinstance(dataset_info_yaml_dict[key]							,				(list, dict, int, str)				)
    _lowercase :							str   	=	yaml.safe_dump(lowerCamelCase_				)
    _lowercase :							str   	=	yaml.safe_load(lowerCamelCase_				)
    assert dataset_info_yaml_dict == reloaded



def        UpperCamelCase_(   )					->		int:
    _lowercase :							Tuple   	=	DatasetInfo()
    _lowercase :							Tuple   	=	dataset_info._to_yaml_dict()
    assert dataset_info_yaml_dict == {}



@pytest.mark.parametrize(
    'dataset_infos_dict'							,				[
        DatasetInfosDict(),
        DatasetInfosDict({'default': DatasetInfo()}				),
        DatasetInfosDict({'my_config_name': DatasetInfo()}				),
        DatasetInfosDict(
            {
                'default': DatasetInfo(
                    description='foo'							,				features=Features({'a': Value('int32'				)}				)							,				builder_name='builder'							,				config_name='config'							,				version='1.0.0'							,				splits=[{'name': 'train'}]							,				download_size=42							,				)
            }				),
        DatasetInfosDict(
            {
                'v1': DatasetInfo(dataset_size=42				),
                'v2': DatasetInfo(dataset_size=1337				),
            }				),
    ]							,				)
def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_				)					->		Dict:
    _lowercase :							Tuple   	=	str(lowerCamelCase_				)
    dataset_infos_dict.write_to_directory(lowerCamelCase_				)
    _lowercase :							Tuple   	=	DatasetInfosDict.from_directory(lowerCamelCase_				)

    # the config_name of the dataset_infos_dict take over the attribute
    for config_name, dataset_info in dataset_infos_dict.items():
        _lowercase :							Any   	=	config_name
        # the yaml representation doesn't include fields like description or citation
        # so we just test that we can recover what we can from the yaml
        _lowercase :							str   	=	DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()				)
    assert dataset_infos_dict == reloaded

    if dataset_infos_dict:
        assert os.path.exists(os.path.join(lowerCamelCase_							,				'README.md'				)				)


