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

import sentencepiece

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


lowerCamelCase__     :		str        = logging.get_logger(__name__)

lowerCamelCase__     :		Optional[int]        = """▁"""

lowerCamelCase__     :		Optional[int]        = {
    """vocab_file""": """vocab.json""",
    """spm_file""": """sentencepiece.bpe.model""",
    """tokenizer_config_file""": """tokenizer_config.json""",
}

lowerCamelCase__     :		int        = {
    """vocab_file""": {
        """facebook/m2m100_418M""": """https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json""",
        """facebook/m2m100_1.2B""": """https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json""",
    },
    """spm_file""": {
        """facebook/m2m100_418M""": """https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model""",
        """facebook/m2m100_1.2B""": """https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model""",
    },
    """tokenizer_config_file""": {
        """facebook/m2m100_418M""": """https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json""",
        """facebook/m2m100_1.2B""": """https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json""",
    },
}

lowerCamelCase__     :		List[Any]        = {
    """facebook/m2m100_418M""": 1_0_2_4,
}

# fmt: off
lowerCamelCase__     :		Any        = {
    """m2m100""": ["""af""", """am""", """ar""", """ast""", """az""", """ba""", """be""", """bg""", """bn""", """br""", """bs""", """ca""", """ceb""", """cs""", """cy""", """da""", """de""", """el""", """en""", """es""", """et""", """fa""", """ff""", """fi""", """fr""", """fy""", """ga""", """gd""", """gl""", """gu""", """ha""", """he""", """hi""", """hr""", """ht""", """hu""", """hy""", """id""", """ig""", """ilo""", """is""", """it""", """ja""", """jv""", """ka""", """kk""", """km""", """kn""", """ko""", """lb""", """lg""", """ln""", """lo""", """lt""", """lv""", """mg""", """mk""", """ml""", """mn""", """mr""", """ms""", """my""", """ne""", """nl""", """no""", """ns""", """oc""", """or""", """pa""", """pl""", """ps""", """pt""", """ro""", """ru""", """sd""", """si""", """sk""", """sl""", """so""", """sq""", """sr""", """ss""", """su""", """sv""", """sw""", """ta""", """th""", """tl""", """tn""", """tr""", """uk""", """ur""", """uz""", """vi""", """wo""", """xh""", """yi""", """yo""", """zh""", """zu"""],
    """wmt21""": ["""en""", """ha""", """is""", """ja""", """cs""", """ru""", """zh""", """de"""]
}



class    _snake_case (						lowerCAmelCase__				):
 __lowerCAmelCase	:  List[str]	   =      VOCAB_FILES_NAMES
 __lowerCAmelCase	:  Tuple	   =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCAmelCase	:  Optional[int]	   =      PRETRAINED_VOCAB_FILES_MAP
 __lowerCAmelCase	:  Union[str, Any]	   =      ['input_ids', 'attention_mask']

 __lowerCAmelCase	:  Dict	   =      []
 __lowerCAmelCase	:  Optional[int]	   =      []


 def __init__(       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_=None       ,	SCREAMING_SNAKE_CASE_=None       ,	SCREAMING_SNAKE_CASE_="<s>"       ,	SCREAMING_SNAKE_CASE_="</s>"       ,	SCREAMING_SNAKE_CASE_="</s>"       ,	SCREAMING_SNAKE_CASE_="<pad>"       ,	SCREAMING_SNAKE_CASE_="<unk>"       ,	SCREAMING_SNAKE_CASE_="m2m100"       ,	SCREAMING_SNAKE_CASE_ = None       ,	SCREAMING_SNAKE_CASE_=8       ,	**SCREAMING_SNAKE_CASE_       ,	):


   '''simple docstring'''




   lowercase__	: Tuple     		=   {} if sp_model_kwargs is None else sp_model_kwargs

   lowercase__	: List[str]     		=   language_codes
   lowercase__	: Optional[int]     		=   FAIRSEQ_LANGUAGE_CODES[language_codes]
   lowercase__	: Tuple     		=   {lang_code: f'__{lang_code}__' for lang_code in fairseq_language_code}

   lowercase__	: str     		=   kwargs.get("""additional_special_tokens"""       ,	[])
   kwargs["additional_special_tokens"] += [
       self.get_lang_token(_lowerCamelCase)
       for lang_code in fairseq_language_code
       if self.get_lang_token(_lowerCamelCase) not in kwargs["additional_special_tokens"]
   ]

   super().__init__(
       src_lang=_lowerCamelCase       ,	tgt_lang=_lowerCamelCase       ,	bos_token=_lowerCamelCase       ,	eos_token=_lowerCamelCase       ,	sep_token=_lowerCamelCase       ,	unk_token=_lowerCamelCase       ,	pad_token=_lowerCamelCase       ,	language_codes=_lowerCamelCase       ,	sp_model_kwargs=self.sp_model_kwargs       ,	num_madeup_words=_lowerCamelCase       ,	**_lowerCamelCase       ,	)

   lowercase__	: Dict     		=   vocab_file
   lowercase__	: str     		=   load_json(_lowerCamelCase)
   lowercase__	: Optional[int]     		=   {v: k for k, v in self.encoder.items()}
   lowercase__	: List[str]     		=   spm_file
   lowercase__	: Any     		=   load_spm(_lowerCamelCase       ,	self.sp_model_kwargs)

   lowercase__	: Tuple     		=   len(self.encoder)

   lowercase__	: int     		=   {
       self.get_lang_token(_lowerCamelCase): self.encoder_size + i for i, lang_code in enumerate(_lowerCamelCase)
   }
   lowercase__	: str     		=   {lang_code: self.encoder_size + i for i, lang_code in enumerate(_lowerCamelCase)}
   lowercase__	: Dict     		=   {v: k for k, v in self.lang_token_to_id.items()}

   lowercase__	: str     		=   src_lang if src_lang is not None else """en"""
   lowercase__	: Any     		=   tgt_lang
   lowercase__	: Union[str, Any]     		=   self.get_lang_id(self._src_lang)
   self.set_src_lang_special_tokens(self._src_lang)

   lowercase__	: Optional[int]     		=   num_madeup_words


 @property
 def   lowercase__  (       self):


   '''simple docstring'''




   return len(self.encoder) + len(self.lang_token_to_id)


 @property
 def   lowercase__  (       self):


   '''simple docstring'''




   return self._src_lang


 @src_lang.setter
 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: Dict     		=   new_src_lang
   self.set_src_lang_special_tokens(self._src_lang)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   if token in self.lang_token_to_id:
     return self.lang_token_to_id[token]
   return self.encoder.get(_lowerCamelCase       ,	self.encoder[self.unk_token])


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   if index in self.id_to_lang_token:
     return self.id_to_lang_token[index]
   return self.decoder.get(_lowerCamelCase       ,	self.unk_token)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: str     		=   []
   lowercase__	: Optional[int]     		=   """"""
   for token in tokens:
     # make sure that special tokens are not decoded using sentencepiece model
     if token in self.all_special_tokens:
       out_string += self.sp_model.decode(_lowerCamelCase) + token
       lowercase__	: List[str]     		=   []
     else:
       current_sub_tokens.append(_lowerCamelCase)
   out_string += self.sp_model.decode(_lowerCamelCase)
   return out_string.strip()


 def   lowercase__  (       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=_lowerCamelCase       ,	token_ids_a=_lowerCamelCase       ,	already_has_special_tokens=_lowerCamelCase)

   lowercase__	: List[Any]     		=   [1] * len(self.prefix_tokens)
   lowercase__	: Union[str, Any]     		=   [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   lowercase__  (       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   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: Any     		=   {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'''




   lowercase__	: Union[str, Any]     		=   self.__dict__.copy()
   lowercase__	: Any     		=   None
   return state


 def __setstate__(       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: int     		=   d

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

   lowercase__	: List[str]     		=   load_spm(self.spm_file       ,	self.sp_model_kwargs)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_ = None):


   '''simple docstring'''




   lowercase__	: Dict     		=   Path(_lowerCamelCase)
   if not save_dir.is_dir():
     raise OSError(f'{save_directory} should be a directory')
   lowercase__	: Any     		=   save_dir / (
       (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""vocab_file"""]
   )
   lowercase__	: Dict     		=   save_dir / (
       (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""spm_file"""]
   )

   save_json(self.encoder       ,	_lowerCamelCase)

   if os.path.abspath(self.spm_file) != os.path.abspath(_lowerCamelCase) and os.path.isfile(self.spm_file):
     copyfile(self.spm_file       ,	_lowerCamelCase)
   elif not os.path.isfile(self.spm_file):
     with open(_lowerCamelCase       ,	"""wb""") as fi:
       lowercase__	: Dict     		=   self.sp_model.serialized_model_proto()
       fi.write(_lowerCamelCase)

   return (str(_lowerCamelCase), str(_lowerCamelCase))


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_ = "en"       ,	SCREAMING_SNAKE_CASE_ = None       ,	SCREAMING_SNAKE_CASE_ = "ro"       ,	**SCREAMING_SNAKE_CASE_       ,	):


   '''simple docstring'''




   lowercase__	: Optional[int]     		=   src_lang
   lowercase__	: Optional[int]     		=   tgt_lang
   self.set_src_lang_special_tokens(self.src_lang)
   return super().prepare_seqaseq_batch(_lowerCamelCase       ,	_lowerCamelCase       ,	**_lowerCamelCase)


 def   lowercase__  (       self       ,	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""")
   lowercase__	: Union[str, Any]     		=   src_lang
   lowercase__	: Union[str, Any]     		=   self(_lowerCamelCase       ,	add_special_tokens=_lowerCamelCase       ,	**_lowerCamelCase)
   lowercase__	: Any     		=   self.get_lang_id(_lowerCamelCase)
   lowercase__	: Any     		=   tgt_lang_id
   return inputs


 def   lowercase__  (       self):


   '''simple docstring'''




   self.set_src_lang_special_tokens(self.src_lang)


 def   lowercase__  (       self):


   '''simple docstring'''




   self.set_tgt_lang_special_tokens(self.tgt_lang)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: List[Any]     		=   self.get_lang_token(_lowerCamelCase)
   lowercase__	: Optional[int]     		=   self.lang_token_to_id[lang_token]
   lowercase__	: Union[str, Any]     		=   [self.cur_lang_id]
   lowercase__	: Optional[int]     		=   [self.eos_token_id]


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: str     		=   self.get_lang_token(_lowerCamelCase)
   lowercase__	: Optional[Any]     		=   self.lang_token_to_id[lang_token]
   lowercase__	: Union[str, Any]     		=   [self.cur_lang_id]
   lowercase__	: Optional[int]     		=   [self.eos_token_id]


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   return self.lang_code_to_token[lang]



 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: Optional[int]     		=   self.get_lang_token(_lowerCamelCase)
   return self.lang_token_to_id[lang_token]

def 		UpperCamelCase						(   lowercase_			,			lowercase_ )		->						sentencepiece.SentencePieceProcessor:


  '''simple docstring'''






  lowercase__	: Dict     		=   sentencepiece.SentencePieceProcessor(**lowercase_ )
  spm.Load(str(lowercase_ ) )
  return spm

def 		UpperCamelCase						(   lowercase_ )		->						Union[Dict, List]:


  '''simple docstring'''






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



def 		UpperCamelCase						(   lowercase_			,			lowercase_ )		->						None:


  '''simple docstring'''






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









"""simple docstring"""
import unittest

import numpy as np
import torch

from .utils_summarization import build_mask, compute_token_type_ids, process_story, truncate_or_pad

class       lowerCamelCase	(       unittest.TestCase   ):


  '''simple docstring'''







  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							List[str]							=    10



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							List[str]							=    [1, 2, 3, 4]
        UpperCAmelCase__			:							Tuple							=    [1, 2, 3, 4, 0, 0, 0, 0, 0, 0]
        self.assertEqual(truncate_or_pad(_lowerCamelCase							,       self.block_size							,       0	)							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Optional[int]							=    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        UpperCAmelCase__			:							Tuple							=    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        self.assertEqual(truncate_or_pad(_lowerCamelCase							,       self.block_size							,       0	)							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Tuple							=    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
        UpperCAmelCase__			:							Dict							=    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        self.assertEqual(truncate_or_pad(_lowerCamelCase							,       self.block_size							,       0	)							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Any							=    """It was the year of Our Lord one thousand seven hundred and
        seventy-five.\n\nSpiritual revelations were conceded to England at that
        favoured period, as at this."""
        UpperCAmelCase__					,		UpperCAmelCase__			:							Any							=    process_story(_lowerCamelCase	)
        self.assertEqual(_lowerCamelCase							,       []	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							str							=    """"""
        UpperCAmelCase__					,		UpperCAmelCase__			:							List[Any]							=    process_story(_lowerCamelCase	)
        self.assertEqual(_lowerCamelCase							,       []	)
        self.assertEqual(_lowerCamelCase							,       []	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Optional[Any]							=    (
            """It was the year of Our Lord one thousand seven hundred and """
            """seventy-five\n\nSpiritual revelations were conceded to England """
            """at that favoured period, as at this.\n@highlight\n\nIt was the best of times"""
        )
        UpperCAmelCase__					,		UpperCAmelCase__			:							Optional[Any]							=    process_story(_lowerCamelCase	)

        UpperCAmelCase__			:							List[Any]							=    [
            """It was the year of Our Lord one thousand seven hundred and seventy-five.""",
            """Spiritual revelations were conceded to England at that favoured period, as at this.""",
        ]
        self.assertEqual(_lowerCamelCase							,       _lowerCamelCase	)

        UpperCAmelCase__			:							List[Any]							=    ["""It was the best of times."""]
        self.assertEqual(_lowerCamelCase							,       _lowerCamelCase	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							str							=    torch.tensor([1, 2, 3, 4]	)
        UpperCAmelCase__			:							str							=    torch.tensor([1, 1, 1, 1]	)
        np.testing.assert_array_equal(build_mask(_lowerCamelCase							,       0	).numpy()							,       expected.numpy()	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							List[str]							=    torch.tensor([1, 2, 3, 4, 23, 23, 23]	)
        UpperCAmelCase__			:							Union[str, Any]							=    torch.tensor([1, 1, 1, 1, 0, 0, 0]	)
        np.testing.assert_array_equal(build_mask(_lowerCamelCase							,       23	).numpy()							,       expected.numpy()	)



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Optional[Any]							=    torch.tensor([8, 2, 3, 4, 1, 1, 1]	)
        UpperCAmelCase__			:							Any							=    torch.tensor([1, 1, 1, 1, 0, 0, 0]	)
        np.testing.assert_array_equal(build_mask(_lowerCamelCase							,       1	).numpy()							,       expected.numpy()	)




  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Any							=    101
        UpperCAmelCase__			:							Dict							=    torch.tensor([[1, 2, 3, 4, 5, 6], [1, 2, 3, 101, 5, 6], [1, 101, 3, 4, 101, 6]]	)
        UpperCAmelCase__			:							Optional[Any]							=    torch.tensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 0, 0], [1, 0, 0, 0, 1, 1]]	)

        UpperCAmelCase__			:							int							=    compute_token_type_ids(_lowerCamelCase							,       _lowerCamelCase	)
        np.testing.assert_array_equal(_lowerCamelCase							,       _lowerCamelCase	)

