

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


lowercase_       		=  logging.get_logger(__name__)

lowercase_       		=  """▁"""

lowercase_       		=  {
    """vocab_file""": """vocab.json""",
    """spm_file""": """sentencepiece.bpe.model""",
    """tokenizer_config_file""": """tokenizer_config.json""",
}

lowercase_       		=  {
    """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""",
    },
}

lowercase_       		=  {
    """facebook/m2m100_418M""": 1024,
}

# fmt: off
lowercase_       		=  {
    """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  A						(							UpperCAmelCase__ ):







    """simple docstring"""





    lowerCamelCase			=							VOCAB_FILES_NAMES
    lowerCamelCase			=							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    lowerCamelCase			=							PRETRAINED_VOCAB_FILES_MAP
    lowerCamelCase			=							['input_ids', 'attention_mask']

    lowerCamelCase			=							[]
    lowerCamelCase			=							[]


    def __init__(      self :       Any,lowercase_ :       Union[str, Any],lowercase_ :       Any,lowercase_ :       Tuple=None,lowercase_ :       List[Any]=None,lowercase_ :       str="<s>",lowercase_ :       Tuple="</s>",lowercase_ :       str="</s>",lowercase_ :       Any="<pad>",lowercase_ :       Union[str, Any]="<unk>",lowercase_ :       int="m2m100",lowercase_ :       Optional[Any] = None,lowercase_ :       Dict=8,**lowercase_ :       int,)->    Optional[int]:





        '''simple docstring'''







        A__          =   {} if sp_model_kwargs is None else sp_model_kwargs

        A__          =   language_codes
        A__          =   FAIRSEQ_LANGUAGE_CODES[language_codes]
        A__          =   {lang_code: F'__{lang_code}__' for lang_code in fairseq_language_code}

        A__          =   kwargs.get('additional_special_tokens',[]		)
        kwargs["additional_special_tokens"] += [
            self.get_lang_token(_a		)
            for lang_code in fairseq_language_code
            if self.get_lang_token(_a		) not in kwargs["additional_special_tokens"]
        ]

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

        A__          =   vocab_file
        A__          =   load_json(_a		)
        A__          =   {v: k for k, v in self.encoder.items()}
        A__          =   spm_file
        A__          =   load_spm(_a,self.sp_model_kwargs		)

        A__          =   len(self.encoder		)

        A__          =   {
            self.get_lang_token(_a		): self.encoder_size + i for i, lang_code in enumerate(_a		)
        }
        A__          =   {lang_code: self.encoder_size + i for i, lang_code in enumerate(_a		)}
        A__          =   {v: k for k, v in self.lang_token_to_id.items()}

        A__          =   src_lang if src_lang is not None else 'en'
        A__          =   tgt_lang
        A__          =   self.get_lang_id(self._src_lang		)
        self.set_src_lang_special_tokens(self._src_lang		)

        A__          =   num_madeup_words


    @property
    def 						snake_case__     (      self :       str		)->    Union[str, Any]:





        '''simple docstring'''







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


    @property
    def 						snake_case__     (      self :       Optional[int]		)->    str:





        '''simple docstring'''







        return self._src_lang


    @src_lang.setter
    def 						snake_case__     (      self :       Tuple,lowercase_ :       Optional[int]		)->    List[Any]:





        '''simple docstring'''







        A__          =   new_src_lang
        self.set_src_lang_special_tokens(self._src_lang		)


    def 						snake_case__     (      self :       int,lowercase_ :       Tuple		)->    List[Any]:





        '''simple docstring'''







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


    def 						snake_case__     (      self :       Any,lowercase_ :       Optional[int]		)->    Tuple:





        '''simple docstring'''







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


    def 						snake_case__     (      self :       List[Any],lowercase_ :       Union[str, Any]		)->    Dict:





        '''simple docstring'''







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


    def 						snake_case__     (      self :       Optional[int],lowercase_ :       Dict		)->    Optional[int]:





        '''simple docstring'''







        A__          =   []
        A__          =   ''
        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(_a		) + token
                A__          =   []
            else:
                current_sub_tokens.append(_a		)
        out_string += self.sp_model.decode(_a		)
        return out_string.strip()


    def 						snake_case__     (      self :       List[str],lowercase_ :       str,lowercase_ :       Any = None,lowercase_ :       str = False		)->    List[str]:





        '''simple docstring'''







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

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


    def 						snake_case__     (      self :       Any,lowercase_ :       Union[str, Any],lowercase_ :       Optional[Any] = None		)->    Tuple:





        '''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 						snake_case__     (      self :       str		)->    Union[str, Any]:





        '''simple docstring'''







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


    def __getstate__(      self :       Any		)->    List[str]:





        '''simple docstring'''







        A__          =   self.__dict__.copy()
        A__          =   None
        return state


    def __setstate__(      self :       int,lowercase_ :       int		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   d

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

        A__          =   load_spm(self.spm_file,self.sp_model_kwargs		)


    def 						snake_case__     (      self :       Optional[int],lowercase_ :       List[Any],lowercase_ :       Union[str, Any] = None		)->    Any:





        '''simple docstring'''







        A__          =   Path(_a		)
        if not save_dir.is_dir():
            raise OSError(F'{save_directory} should be a directory'		)
        A__          =   save_dir / (
            (filename_prefix + '-' if filename_prefix else '') + self.vocab_files_names['vocab_file']
        )
        A__          =   save_dir / (
            (filename_prefix + '-' if filename_prefix else '') + self.vocab_files_names['spm_file']
        )

        save_json(self.encoder,_a		)

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

        return (str(_a		), str(_a		))


    def 						snake_case__     (      self :       Any,lowercase_ :       Union[str, Any],lowercase_ :       str = "en",lowercase_ :       Union[str, Any] = None,lowercase_ :       Optional[Any] = "ro",**lowercase_ :       List[str],)->    Tuple:





        '''simple docstring'''







        A__          =   src_lang
        A__          =   tgt_lang
        self.set_src_lang_special_tokens(self.src_lang		)
        return super().prepare_seqaseq_batch(_a,_a,**_a		)


    def 						snake_case__     (      self :       int,lowercase_ :       int,lowercase_ :       Optional[int],lowercase_ :       int,**lowercase_ :       Tuple		)->    int:





        '''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'		)
        A__          =   src_lang
        A__          =   self(_a,add_special_tokens=_a,**_a		)
        A__          =   self.get_lang_id(_a		)
        A__          =   tgt_lang_id
        return inputs


    def 						snake_case__     (      self :       int		)->    Optional[int]:





        '''simple docstring'''







        self.set_src_lang_special_tokens(self.src_lang		)


    def 						snake_case__     (      self :       List[Any]		)->    int:





        '''simple docstring'''







        self.set_tgt_lang_special_tokens(self.tgt_lang		)


    def 						snake_case__     (      self :       List[Any],lowercase_ :       Union[str, Any]		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   self.get_lang_token(_a		)
        A__          =   self.lang_token_to_id[lang_token]
        A__          =   [self.cur_lang_id]
        A__          =   [self.eos_token_id]


    def 						snake_case__     (      self :       Tuple,lowercase_ :       str		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   self.get_lang_token(_a		)
        A__          =   self.lang_token_to_id[lang_token]
        A__          =   [self.cur_lang_id]
        A__          =   [self.eos_token_id]


    def 						snake_case__     (      self :       Tuple,lowercase_ :       Union[str, Any]		)->    Optional[Any]:





        '''simple docstring'''







        return self.lang_code_to_token[lang]


    def 						snake_case__     (      self :       str,lowercase_ :       List[Any]		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   self.get_lang_token(_a		)
        return self.lang_token_to_id[lang_token]




def 						_snake_case(    SCREAMING_SNAKE_CASE__			:	Dict    ,			SCREAMING_SNAKE_CASE__			:	Dict      )							->      sentencepiece.SentencePieceProcessor:





    '''simple docstring'''






    A__          =   sentencepiece.SentencePieceProcessor(**snake_case__      )
    spm.Load(str(snake_case__      )      )
    return spm
def 						_snake_case(    SCREAMING_SNAKE_CASE__			:	Dict      )							->      Union[Dict, List]:





    '''simple docstring'''






    with open(snake_case__    ,			'r'      ) as f:
        return json.load(snake_case__      )

def 						_snake_case(    SCREAMING_SNAKE_CASE__			:	Optional[Any]    ,			SCREAMING_SNAKE_CASE__			:	Any      )							->      None:





    '''simple docstring'''






    with open(snake_case__    ,			'w'      ) as f:
        json.dump(snake_case__    ,			snake_case__    ,			indent=2      )







"""simple docstring"""




from math import ceil



def     a__ (				snake_case__       , snake_case__				)						->					Optional[int]:
			lowerCamelCase			      =					list(range(0       , snake_case__				)				)

			lowerCamelCase			      =					[item for sublist in list(device_map.values()				) for item in sublist]

			# Duplicate check
			lowerCamelCase			      =					[]
			for i in device_map_blocks:
						if device_map_blocks.count(snake_case__				) > 1 and i not in duplicate_blocks:
									duplicate_blocks.append(snake_case__				)
    # Missing blocks
			lowerCamelCase			      =					[i for i in blocks if i not in device_map_blocks]
			lowerCamelCase			      =					[i for i in device_map_blocks if i not in blocks]

			if len(snake_case__				) != 0:
						raise ValueError(
						    """Duplicate attention blocks specified in device_map. Attention blocks must be specified to one device."""
						    """ These attention blocks were specified more than once: """ + str(snake_case__				)				)
			if len(snake_case__				) != 0:
						raise ValueError(
						    """There are attention blocks for this model that are not specified in the device_map. Add these attention """
						    """blocks to a device on the device_map: """ + str(snake_case__				)				)
			if len(snake_case__				) != 0:
						raise ValueError(
						    """The device_map contains more attention blocks than this model has. Remove these from the device_map:"""
						    + str(snake_case__				)				)



def     a__ (				snake_case__       , snake_case__				)						->					List[Any]:
			lowerCamelCase			      =					list(range(snake_case__				)				)
			lowerCamelCase			      =					int(ceil(n_layers / len(snake_case__				)				)				)
			lowerCamelCase			      =					[layers[i : i + n_blocks] for i in range(0       , snake_case__       , snake_case__				)]

			return dict(zip(snake_case__       , snake_case__				)				)



