




import json
import os
import re
import warnings
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowercase_     :    Optional[int]      =			logging.get_logger(__name__)

lowercase_     :    List[Any]      =			{
    'source_spm': 'source.spm',
    'target_spm': 'target.spm',
    'vocab': 'vocab.json',
    'target_vocab_file': 'target_vocab.json',
    'tokenizer_config_file': 'tokenizer_config.json',
}

lowercase_     :    List[str]      =			{
    'source_spm': {
        'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/source.spm'
    },
    'target_spm': {
        'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/target.spm'
    },
    'vocab': {
        'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/vocab.json'
    },
    'tokenizer_config_file': {
        'Helsinki-NLP/opus-mt-en-de': (
            'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/tokenizer_config.json'
        )
    },
}

lowercase_     :    List[str]      =			{'Helsinki-NLP/opus-mt-en-de': 5_12}
lowercase_     :    str      =			{}






class    __lowerCAmelCase			(	UpperCAmelCase__			):
			snake_case_				:			Tuple								=					VOCAB_FILES_NAMES
			snake_case_				:			Union[str, Any]								=					PRETRAINED_VOCAB_FILES_MAP
			snake_case_				:			Union[str, Any]								=					PRETRAINED_INIT_CONFIGURATION
			snake_case_				:			Dict								=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			snake_case_				:			str								=					["input_ids", "attention_mask"]
			snake_case_				:			int								=					re.compile(">>.+<<"			)  # type: re.Pattern







			def __init__(	self   :      List[str]   ,   snake_case__   :      Optional[int]   ,   snake_case__   :      str   ,   snake_case__   :      str   ,   snake_case__   :      str=None   ,   snake_case__   :      str=None   ,   snake_case__   :      Union[str, Any]=None   ,   snake_case__   :      Dict="<unk>"   ,   snake_case__   :      Optional[int]="</s>"   ,   snake_case__   :      Optional[int]="<pad>"   ,   snake_case__   :      Optional[Any]=512   ,   snake_case__   :      Optional[Dict[str, Any]] = None   ,   snake_case__   :      Tuple=False   ,   **snake_case__   :      str   ,   ):




									"""simple docstring"""


									_UpperCAmelCase													=					{} if sp_model_kwargs is None else sp_model_kwargs

									super().__init__(
									    # bos_token=bos_token,  unused. Start decoding with config.decoder_start_token_id
									    source_lang=snake_case__   ,   target_lang=snake_case__   ,   unk_token=snake_case__   ,   eos_token=snake_case__   ,   pad_token=snake_case__   ,   model_max_length=snake_case__   ,   sp_model_kwargs=self.sp_model_kwargs   ,   target_vocab_file=snake_case__   ,   separate_vocabs=snake_case__   ,   **snake_case__   ,   )
									assert Path(snake_case__				).exists(), F"""cannot find spm source {source_spm}"""

									_UpperCAmelCase													=					separate_vocabs
									_UpperCAmelCase													=					load_json(snake_case__				)
									if self.unk_token not in self.encoder:
															raise KeyError("<unk> token must be in vocab"				)
									assert self.pad_token in self.encoder

									if separate_vocabs:
															_UpperCAmelCase													=					load_json(snake_case__				)
															_UpperCAmelCase													=					{v: k for k, v in self.target_encoder.items()}
															_UpperCAmelCase													=					[]
									else:
															_UpperCAmelCase													=					{v: k for k, v in self.encoder.items()}
															_UpperCAmelCase													=					[k for k in self.encoder if k.startswith(">>"				) and k.endswith("<<"				)]

									_UpperCAmelCase													=					source_lang
									_UpperCAmelCase													=					target_lang
									_UpperCAmelCase													=					[source_spm, target_spm]

									# load SentencePiece model for pre-processing
									_UpperCAmelCase													=					load_spm(snake_case__   ,   self.sp_model_kwargs				)
									_UpperCAmelCase													=					load_spm(snake_case__   ,   self.sp_model_kwargs				)
									_UpperCAmelCase													=					self.spm_source
									_UpperCAmelCase													=					self.encoder

									# Multilingual target side: default to using first supported language code.

									self._setup_normalizer()







			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									try:
															from sacremoses import MosesPunctNormalizer

															_UpperCAmelCase													=					MosesPunctNormalizer(self.source_lang				).normalize
									except (ImportError, FileNotFoundError):
															warnings.warn("Recommended: pip install sacremoses."				)
															_UpperCAmelCase													=					lambda snake_case__				: x







			def UpperCamelCase      (	self   :      List[Any]   ,   snake_case__   :      str				):




									"""simple docstring"""


									return self.punc_normalizer(snake_case__				) if x else ""







			def UpperCamelCase      (	self   :      Union[str, Any]   ,   snake_case__   :      Optional[int]				):




									"""simple docstring"""


									return self.current_encoder.get(snake_case__   ,   self.current_encoder[self.unk_token]				)







			def UpperCamelCase      (	self   :      int   ,   snake_case__   :      str				):




									"""simple docstring"""


									_UpperCAmelCase													=					self.language_code_re.match(snake_case__				)
									_UpperCAmelCase													=					[match.group(0				)] if match else []
									return code, self.language_code_re.sub(""   ,   snake_case__				)







			def UpperCamelCase      (	self   :      List[Any]   ,   snake_case__   :      str				):




									"""simple docstring"""


									_UpperCAmelCase							,			_UpperCAmelCase													=					self.remove_language_code(snake_case__				)
									_UpperCAmelCase													=					self.current_spm.encode(snake_case__   ,   out_type=snake_case__				)
									return code + pieces







			def UpperCamelCase      (	self   :      Dict   ,   snake_case__   :      int				):




									"""simple docstring"""


									return self.decoder.get(snake_case__   ,   self.unk_token				)







			def UpperCamelCase      (	self   :      int   ,   snake_case__   :      int   ,   **snake_case__   :      List[str]				):




									"""simple docstring"""


									return super().batch_decode(snake_case__   ,   **snake_case__				)







			def UpperCamelCase      (	self   :      int   ,   snake_case__   :      List[Any]   ,   **snake_case__   :      List[Any]				):




									"""simple docstring"""


									return super().decode(snake_case__   ,   **snake_case__				)







			def UpperCamelCase      (	self   :      Optional[Any]   ,   snake_case__   :      List[str]				):




									"""simple docstring"""


									_UpperCAmelCase													=					self.spm_source if self._decode_use_source_tokenizer else self.spm_target
									_UpperCAmelCase													=					[]
									_UpperCAmelCase													=					""
									for token in tokens:
															# make sure that special tokens are not decoded using sentencepiece model
															if token in self.all_special_tokens:
																					out_string += sp_model.decode_pieces(snake_case__				) + token + " "
																					_UpperCAmelCase													=					[]
															else:
																					current_sub_tokens.append(snake_case__				)
									out_string += sp_model.decode_pieces(snake_case__				)
									return out_string.strip()







			def UpperCamelCase      (	self   :      Dict   ,   snake_case__   :      Tuple   ,   snake_case__   :      Tuple=None				):




									"""simple docstring"""


									if token_ids_a is None:
															return token_ids_a + [self.eos_token_id]
									# We don't expect to process pairs, but leave the pair logic for API consistency
									return token_ids_a + token_ids_a + [self.eos_token_id]







			def UpperCamelCase      (	self   :      Any				):




									"""simple docstring"""


									_UpperCAmelCase													=					self.spm_source
									_UpperCAmelCase													=					self.encoder







			def UpperCamelCase      (	self   :      str				):




									"""simple docstring"""


									_UpperCAmelCase													=					self.spm_target
									if self.separate_vocabs:
															_UpperCAmelCase													=					self.target_encoder







			@property
			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									return len(self.encoder				)







			def UpperCamelCase      (	self   :      str   ,   snake_case__   :      str   ,   snake_case__   :      Optional[str] = None				):




									"""simple docstring"""


									if not os.path.isdir(snake_case__				):
															logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""				)
															return
									_UpperCAmelCase													=					[]

									if self.separate_vocabs:
															_UpperCAmelCase													=					os.path.join(
															    snake_case__   ,   (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab"]   ,   )
															_UpperCAmelCase													=					os.path.join(
															    snake_case__   ,   (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["target_vocab_file"]   ,   )
															save_json(self.encoder   ,   snake_case__				)
															save_json(self.target_encoder   ,   snake_case__				)
															saved_files.append(snake_case__				)
															saved_files.append(snake_case__				)
									else:
															_UpperCAmelCase													=					os.path.join(
															    snake_case__   ,   (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab"]				)
															save_json(self.encoder   ,   snake_case__				)
															saved_files.append(snake_case__				)

									for spm_save_filename, spm_orig_path, spm_model in zip(
									    [VOCAB_FILES_NAMES["source_spm"], VOCAB_FILES_NAMES["target_spm"]]   ,   self.spm_files   ,   [self.spm_source, self.spm_target]   ,   ):
															_UpperCAmelCase													=					os.path.join(
															    snake_case__   ,   (filename_prefix + "-" if filename_prefix else "") + spm_save_filename				)
															if os.path.abspath(snake_case__				) != os.path.abspath(snake_case__				) and os.path.isfile(snake_case__				):
																					copyfile(snake_case__   ,   snake_case__				)
																					saved_files.append(snake_case__				)
															elif not os.path.isfile(snake_case__				):
																					with open(snake_case__   ,   "wb"				) as fi:
																											_UpperCAmelCase													=					spm_model.serialized_model_proto()
																											fi.write(snake_case__				)
																					saved_files.append(snake_case__				)

									return tuple(snake_case__				)







			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									return self.get_src_vocab()







			def UpperCamelCase      (	self   :      List[str]				):




									"""simple docstring"""


									return dict(self.encoder   ,   **self.added_tokens_encoder				)







			def UpperCamelCase      (	self   :      List[Any]				):




									"""simple docstring"""


									return dict(self.target_encoder   ,   **self.added_tokens_decoder				)







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




									"""simple docstring"""


									_UpperCAmelCase													=					self.__dict__.copy()
									state.update(
									    {k: None for k in ["spm_source", "spm_target", "current_spm", "punc_normalizer", "target_vocab_file"]}				)
									return state







			def __setstate__(	self   :      str   ,   snake_case__   :      Dict				):




									"""simple docstring"""


									_UpperCAmelCase													=					d

									# for backward compatibility
									if not hasattr(self   ,   "sp_model_kwargs"				):
															_UpperCAmelCase													=					{}

									_UpperCAmelCase							,			_UpperCAmelCase													=					(load_spm(snake_case__   ,   self.sp_model_kwargs				) for f in self.spm_files)
									_UpperCAmelCase													=					self.spm_source
									self._setup_normalizer()







			def UpperCamelCase      (	self   :      str   ,   *snake_case__   :      Optional[int]   ,   **snake_case__   :      Union[str, Any]				):




									"""simple docstring"""


									return 1







			def UpperCamelCase      (	self   :      List[Any]   ,   snake_case__   :      Dict				):




									"""simple docstring"""


									_UpperCAmelCase													=					set(self.all_special_ids				)  # call it once instead of inside list comp
									all_special_ids.remove(self.unk_token_id				)  # <unk> is only sometimes special
									return [1 if x in all_special_ids else 0 for x in seq]







			def UpperCamelCase      (	self   :      Tuple   ,   snake_case__   :      List   ,   snake_case__   :      Optional[List] = None   ,   snake_case__   :      bool = False				):




									"""simple docstring"""


									if already_has_special_tokens:
															return self._special_token_mask(snake_case__				)
									elif token_ids_a is None:
															return self._special_token_mask(snake_case__				) + [1]
									else:
															return self._special_token_mask(token_ids_a + token_ids_a				) + [1]




def 							__SCREAMING_SNAKE_CASE  (      snake_case_     ,   snake_case_						):

						'''simple docstring'''


						_UpperCAmelCase													=					sentencepiece.SentencePieceProcessor(**snake_case_						)
						spm.Load(snake_case_						)
						return spm


def 							__SCREAMING_SNAKE_CASE  (      snake_case_     ,   snake_case_						):

						'''simple docstring'''


						with open(snake_case_     ,   "w"						) as f:
												json.dump(snake_case_     ,   snake_case_     ,   indent=2						)





def 							__SCREAMING_SNAKE_CASE  (      snake_case_						):

						'''simple docstring'''


						with open(snake_case_     ,   "r"						) as f:
												return json.load(snake_case_						)










import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


lowercase_     :    Union[str, Any]      =			datasets.utils.logging.get_logger(__name__)






@dataclass
class    __lowerCAmelCase			(	datasets.BuilderConfig			):
			snake_case_				:			Optional[datasets.Features]								=					None






class    __lowerCAmelCase			(	datasets.ArrowBasedBuilder			):
			snake_case_				:			str								=					ArrowConfig







			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									return datasets.DatasetInfo(features=self.config.features				)







			def UpperCamelCase      (	self   :      str   ,   snake_case__   :      Optional[int]				):




									"""simple docstring"""


									if not self.config.data_files:
															raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""				)
									_UpperCAmelCase													=					dl_manager.download_and_extract(self.config.data_files				)
									if isinstance(snake_case__   ,   (str, list, tuple)				):
															_UpperCAmelCase													=					data_files
															if isinstance(snake_case__   ,   snake_case__				):
																					_UpperCAmelCase													=					[files]
															# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
															_UpperCAmelCase													=					[dl_manager.iter_files(snake_case__				) for file in files]
															return [datasets.SplitGenerator(name=datasets.Split.TRAIN   ,   gen_kwargs={"files": files}				)]
									_UpperCAmelCase													=					[]
									for split_name, files in data_files.items():
															if isinstance(snake_case__   ,   snake_case__				):
																					_UpperCAmelCase													=					[files]
															# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
															_UpperCAmelCase													=					[dl_manager.iter_files(snake_case__				) for file in files]
															# Infer features is they are stoed in the arrow schema
															if self.info.features is None:
																					for file in itertools.chain.from_iterable(snake_case__				):
																											with open(snake_case__   ,   "rb"				) as f:
																																	_UpperCAmelCase													=					datasets.Features.from_arrow_schema(pa.ipc.open_stream(snake_case__				).schema				)
																											break
															splits.append(datasets.SplitGenerator(name=snake_case__   ,   gen_kwargs={"files": files}				)				)
									return splits







			def UpperCamelCase      (	self   :      int   ,   snake_case__   :      pa.Table				):




									"""simple docstring"""


									if self.info.features is not None:
															# more expensive cast to support nested features with keys in a different order
															# allows str <-> int/float or str to Audio for example
															_UpperCAmelCase													=					table_cast(snake_case__   ,   self.info.features.arrow_schema				)
									return pa_table







			def UpperCamelCase      (	self   :      Optional[Any]   ,   snake_case__   :      int				):




									"""simple docstring"""


									for file_idx, file in enumerate(itertools.chain.from_iterable(snake_case__				)				):
															with open(snake_case__   ,   "rb"				) as f:
																					try:
																											for batch_idx, record_batch in enumerate(pa.ipc.open_stream(snake_case__				)				):
																																	_UpperCAmelCase													=					pa.Table.from_batches([record_batch]				)
																																	# Uncomment for debugging (will print the Arrow table size and elements)
																																	# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																																	# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																																	yield F"""{file_idx}_{batch_idx}""", self._cast_table(snake_case__				)
																					except ValueError as e:
																											logger.error(F"""Failed to read file '{file}' with error {type(snake_case__				)}: {e}"""				)
																											raise




