

'''simple docstring'''



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


_UpperCamelCase						:  List[Any]							       =				logging.get_logger(__name__)

_UpperCamelCase						:  List[str]							       =				"▁"

_UpperCamelCase						:  Tuple							       =				{"vocab_file": "sentencepiece.bpe.model"}

_UpperCamelCase						:  int							       =				{
    "vocab_file": {
        "facebook/mbart-large-en-ro": (
            "https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/sentencepiece.bpe.model"
        ),
        "facebook/mbart-large-cc25": (
            "https://huggingface.co/facebook/mbart-large-cc25/resolve/main/sentencepiece.bpe.model"
        ),
    }
}

_UpperCamelCase						:  str							       =				{
    "facebook/mbart-large-en-ro": 1024,
    "facebook/mbart-large-cc25": 1024,
}

# fmt: off
_UpperCamelCase						:  Any							       =				["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"]



class 				_snake_case					(   a_			):
  SCREAMING_SNAKE_CASE			:       str				=     VOCAB_FILES_NAMES
  SCREAMING_SNAKE_CASE			:       Any				=     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  SCREAMING_SNAKE_CASE			:       Optional[Any]				=     PRETRAINED_VOCAB_FILES_MAP
  SCREAMING_SNAKE_CASE			:       Optional[Any]				=     ['''input_ids''', '''attention_mask''']

  SCREAMING_SNAKE_CASE			:       List[int]				=     []
  SCREAMING_SNAKE_CASE			:       List[int]				=     []

  def __init__(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE="<s>"      , _SCREAMING_SNAKE_CASE="</s>"      , _SCREAMING_SNAKE_CASE="</s>"      , _SCREAMING_SNAKE_CASE="<s>"      , _SCREAMING_SNAKE_CASE="<unk>"      , _SCREAMING_SNAKE_CASE="<pad>"      , _SCREAMING_SNAKE_CASE="<mask>"      , _SCREAMING_SNAKE_CASE=None      , _SCREAMING_SNAKE_CASE=None      , _SCREAMING_SNAKE_CASE=None      , _SCREAMING_SNAKE_CASE = None      , _SCREAMING_SNAKE_CASE=None      , **_SCREAMING_SNAKE_CASE      , ):





      '''simple docstring'''

      lowerCAmelCase 						=							AddedToken(_SCREAMING_SNAKE_CASE      , lstrip=_SCREAMING_SNAKE_CASE      , rstrip=_SCREAMING_SNAKE_CASE	) if isinstance(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE	) else mask_token

      lowerCAmelCase 						=							{} if sp_model_kwargs is None else sp_model_kwargs

      super().__init__(
          bos_token=_SCREAMING_SNAKE_CASE      , eos_token=_SCREAMING_SNAKE_CASE      , unk_token=_SCREAMING_SNAKE_CASE      , sep_token=_SCREAMING_SNAKE_CASE      , cls_token=_SCREAMING_SNAKE_CASE      , pad_token=_SCREAMING_SNAKE_CASE      , mask_token=_SCREAMING_SNAKE_CASE      , tokenizer_file=_SCREAMING_SNAKE_CASE      , src_lang=_SCREAMING_SNAKE_CASE      , tgt_lang=_SCREAMING_SNAKE_CASE      , additional_special_tokens=_SCREAMING_SNAKE_CASE      , sp_model_kwargs=self.sp_model_kwargs      , **_SCREAMING_SNAKE_CASE      , )

      lowerCAmelCase 						=							spm.SentencePieceProcessor(**self.sp_model_kwargs	)
      self.sp_model.Load(str(_SCREAMING_SNAKE_CASE	)	)
      lowerCAmelCase 						=							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
      lowerCAmelCase 						=							{'<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
      lowerCAmelCase 						=							1

      lowerCAmelCase 						=							len(self.sp_model	)
      lowerCAmelCase 						=							{
          code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(_SCREAMING_SNAKE_CASE	)
      }
      lowerCAmelCase 						=							{v: k for k, v in self.lang_code_to_id.items()}
      lowerCAmelCase 						=							len(self.sp_model	) + len(self.lang_code_to_id	) + self.fairseq_offset

      self.fairseq_tokens_to_ids.update(self.lang_code_to_id	)
      lowerCAmelCase 						=							{v: k for k, v in self.fairseq_tokens_to_ids.items()}
      lowerCAmelCase 						=							list(self.lang_code_to_id.keys()	)

      if additional_special_tokens is not None:
          # Only add those special tokens if they are not already there.
          self._additional_special_tokens.extend(
              [t for t in additional_special_tokens if t not in self._additional_special_tokens]	)

      lowerCAmelCase 						=							src_lang if src_lang is not None else 'en_XX'
      lowerCAmelCase 						=							self.lang_code_to_id[self._src_lang]
      lowerCAmelCase 						=							tgt_lang
      self.set_src_lang_special_tokens(self._src_lang	)

  def __getstate__(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.__dict__.copy()
      lowerCAmelCase 						=							None
      lowerCAmelCase 						=							self.sp_model.serialized_model_proto()
      return state

  def __setstate__(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							d

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

      lowerCAmelCase 						=							spm.SentencePieceProcessor(**self.sp_model_kwargs	)
      self.sp_model.LoadFromSerializedProto(self.sp_model_proto	)

  @property
  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''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   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      return self._src_lang

  @src_lang.setter
  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							new_src_lang
      self.set_src_lang_special_tokens(self._src_lang	)

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE = None      , _SCREAMING_SNAKE_CASE = False	):





      '''simple docstring'''

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

      lowerCAmelCase 						=							[1] * len(self.prefix_tokens	)
      lowerCAmelCase 						=							[1] * len(self.suffix_tokens	)
      if token_ids_a is None:
          return prefix_ones + ([0] * len(_SCREAMING_SNAKE_CASE	)) + suffix_ones
      return prefix_ones + ([0] * len(_SCREAMING_SNAKE_CASE	)) + ([0] * len(_SCREAMING_SNAKE_CASE	)) + suffix_ones

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE = None	):





      '''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   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE = None	):





      '''simple docstring'''

      lowerCAmelCase 						=							[self.sep_token_id]
      lowerCAmelCase 						=							[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]

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	):





      '''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'	)
      lowerCAmelCase 						=							src_lang
      lowerCAmelCase 						=							self(_SCREAMING_SNAKE_CASE      , add_special_tokens=_SCREAMING_SNAKE_CASE      , return_tensors=_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							self.convert_tokens_to_ids(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							tgt_lang_id
      return inputs

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							{self.convert_ids_to_tokens(_SCREAMING_SNAKE_CASE	): i for i in range(self.vocab_size	)}
      vocab.update(self.added_tokens_encoder	)
      return vocab

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      return self.sp_model.encode(_SCREAMING_SNAKE_CASE      , out_type=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      if token in self.fairseq_tokens_to_ids:
          return self.fairseq_tokens_to_ids[token]
      lowerCAmelCase 						=							self.sp_model.PieceToId(_SCREAMING_SNAKE_CASE	)

      # 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   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''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   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							''.join(_SCREAMING_SNAKE_CASE	).replace(_SCREAMING_SNAKE_CASE      , ' '	).strip()
      return out_string

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE = None	):





      '''simple docstring'''

      if not os.path.isdir(_SCREAMING_SNAKE_CASE	):
          logger.error(F'Vocabulary path ({save_directory}) should be a directory'	)
          return
      lowerCAmelCase 						=							os.path.join(
          _SCREAMING_SNAKE_CASE      , (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']	)

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

      return (out_vocab_file,)

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE = "en_XX"      , _SCREAMING_SNAKE_CASE = None      , _SCREAMING_SNAKE_CASE = "ro_RO"      , **_SCREAMING_SNAKE_CASE      , ):





      '''simple docstring'''

      lowerCAmelCase 						=							src_lang
      lowerCAmelCase 						=							tgt_lang
      return super().prepare_seqaseq_batch(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      return self.set_src_lang_special_tokens(self.src_lang	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      return self.set_tgt_lang_special_tokens(self.tgt_lang	)

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.lang_code_to_id[src_lang]
      lowerCAmelCase 						=							[]
      lowerCAmelCase 						=							[self.eos_token_id, self.cur_lang_code]


  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.lang_code_to_id[lang]
      lowerCAmelCase 						=							[]
      lowerCAmelCase 						=							[self.eos_token_id, self.cur_lang_code]








'''simple docstring'''



import itertools
from dataclasses import dataclass
from typing import Optional

import pandas as pd
import pyarrow as pa

import datasets
from datasets.table import table_cast



@dataclass
class 				_snake_case					(   datasets.BuilderConfig			):
  SCREAMING_SNAKE_CASE			:       Optional[datasets.Features]				=     None



class 				_snake_case					(   datasets.ArrowBasedBuilder			):
  SCREAMING_SNAKE_CASE			:       Tuple				=     PandasConfig

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

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

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''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}'	)
      lowerCAmelCase 						=							dl_manager.download_and_extract(self.config.data_files	)
      if isinstance(_SCREAMING_SNAKE_CASE      , (str, list, tuple)	):
          lowerCAmelCase 						=							data_files
          if isinstance(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE	):
              lowerCAmelCase 						=							[files]
          # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
          lowerCAmelCase 						=							[dl_manager.iter_files(_SCREAMING_SNAKE_CASE	) for file in files]
          return [datasets.SplitGenerator(name=datasets.Split.TRAIN      , gen_kwargs={'files': files}	)]
      lowerCAmelCase 						=							[]
      for split_name, files in data_files.items():
          if isinstance(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE	):
              lowerCAmelCase 						=							[files]
          # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
          lowerCAmelCase 						=							[dl_manager.iter_files(_SCREAMING_SNAKE_CASE	) for file in files]
          splits.append(datasets.SplitGenerator(name=_SCREAMING_SNAKE_CASE      , gen_kwargs={'files': files}	)	)
      return splits

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      if self.config.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
          lowerCAmelCase 						=							table_cast(_SCREAMING_SNAKE_CASE      , self.config.features.arrow_schema	)
      return pa_table


  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      for i, file in enumerate(itertools.chain.from_iterable(_SCREAMING_SNAKE_CASE	)	):
          with open(_SCREAMING_SNAKE_CASE      , 'rb'	) as f:
              lowerCAmelCase 						=							pa.Table.from_pandas(pd.read_pickle(_SCREAMING_SNAKE_CASE	)	)
              yield i, self._cast_table(_SCREAMING_SNAKE_CASE	)





