

import json
import os
import re
import unicodedata
from typing import Dict, List, Optional, Tuple

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


lowerCamelCase									= logging.get_logger(__name__)

lowerCamelCase									= {
    '''src_vocab_file''': '''vocab-src.json''',
    '''tgt_vocab_file''': '''vocab-tgt.json''',
    '''merges_file''': '''merges.txt''',
}

lowerCamelCase									= {
    '''src_vocab_file''': {
        '''stas/tiny-wmt19-en-de''': '''https://huggingface.co/stas/tiny-wmt19-en-de/resolve/main/vocab-src.json'''
    },
    '''tgt_vocab_file''': {
        '''stas/tiny-wmt19-en-de''': '''https://huggingface.co/stas/tiny-wmt19-en-de/resolve/main/vocab-tgt.json'''
    },
    '''merges_file''': {'''stas/tiny-wmt19-en-de''': '''https://huggingface.co/stas/tiny-wmt19-en-de/resolve/main/merges.txt'''},
}

lowerCamelCase									= {'''stas/tiny-wmt19-en-de''': 1024}
lowerCamelCase									= {
    '''stas/tiny-wmt19-en-de''': {
        '''langs''': ['''en''', '''de'''],
        '''model_max_length''': 1024,
        '''special_tokens_map_file''': None,
        '''full_tokenizer_file''': None,
    }
}




def       UpperCAmelCase__    (      _A				:			Any							):





  '''simple docstring'''







  a__  =set()
  a__  =word[0]
  for char in word[1:]:
    pairs.add((prev_char, char)							)
    a__  =char
  return pairs




def       UpperCAmelCase__    (      _A				:			str							):





  '''simple docstring'''







  a__  =text.replace('''，'''     ,				''','''							)
  a__  =re.sub(r'''。\s*'''     ,				'''. '''     ,				_A							)
  a__  =text.replace('''、'''     ,				''','''							)
  a__  =text.replace('''”'''     ,				'''"'''							)
  a__  =text.replace('''“'''     ,				'''"'''							)
  a__  =text.replace('''∶'''     ,				''':'''							)
  a__  =text.replace('''：'''     ,				''':'''							)
  a__  =text.replace('''？'''     ,				'''?'''							)
  a__  =text.replace('''《'''     ,				'''"'''							)
  a__  =text.replace('''》'''     ,				'''"'''							)
  a__  =text.replace('''）'''     ,				''')'''							)
  a__  =text.replace('''！'''     ,				'''!'''							)
  a__  =text.replace('''（'''     ,				'''('''							)
  a__  =text.replace('''；'''     ,				''';'''							)
  a__  =text.replace('''１'''     ,				'''1'''							)
  a__  =text.replace('''」'''     ,				'''"'''							)
  a__  =text.replace('''「'''     ,				'''"'''							)
  a__  =text.replace('''０'''     ,				'''0'''							)
  a__  =text.replace('''３'''     ,				'''3'''							)
  a__  =text.replace('''２'''     ,				'''2'''							)
  a__  =text.replace('''５'''     ,				'''5'''							)
  a__  =text.replace('''６'''     ,				'''6'''							)
  a__  =text.replace('''９'''     ,				'''9'''							)
  a__  =text.replace('''７'''     ,				'''7'''							)
  a__  =text.replace('''８'''     ,				'''8'''							)
  a__  =text.replace('''４'''     ,				'''4'''							)
  a__  =re.sub(r'''．\s*'''     ,				'''. '''     ,				_A							)
  a__  =text.replace('''～'''     ,				'''~'''							)
  a__  =text.replace('''’'''     ,				'''\''''							)
  a__  =text.replace('''…'''     ,				'''...'''							)
  a__  =text.replace('''━'''     ,				'''-'''							)
  a__  =text.replace('''〈'''     ,				'''<'''							)
  a__  =text.replace('''〉'''     ,				'''>'''							)
  a__  =text.replace('''【'''     ,				'''['''							)
  a__  =text.replace('''】'''     ,				''']'''							)
  a__  =text.replace('''％'''     ,				'''%'''							)
  return text




def       UpperCAmelCase__    (      _A				:			Optional[int]							):





  '''simple docstring'''







  a__  =[]
  for char in text:
    a__  =unicodedata.category(_A							)
    if cat.startswith('''C'''							):
      continue
    output.append(_A							)
  return "".join(_A							)







class 						__magic_name__      (							lowerCamelCase__							):

   '''simple docstring'''




   lowerCamelCase__			:     Union[str, Any]    =       VOCAB_FILES_NAMES
   lowerCamelCase__			:     List[str]    =       PRETRAINED_VOCAB_FILES_MAP
   lowerCamelCase__			:     List[str]    =       PRETRAINED_INIT_CONFIGURATION
   lowerCamelCase__			:     List[Any]    =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   lowerCamelCase__			:     Union[str, Any]    =       ['input_ids', 'attention_mask']



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






     """simple docstring"""






     super().__init__(
         langs=lowercase_,   src_vocab_file=lowercase_,   tgt_vocab_file=lowercase_,   merges_file=lowercase_,   do_lower_case=lowercase_,   unk_token=lowercase_,   bos_token=lowercase_,   sep_token=lowercase_,   pad_token=lowercase_,   **lowercase_,   )

     try:
       import sacremoses
     except ImportError:
       raise ImportError(
           '''You need to install sacremoses to use XLMTokenizer. '''
           '''See https://pypi.org/project/sacremoses/ for installation.'''	)

     a__  =sacremoses

     a__  =src_vocab_file
     a__  =tgt_vocab_file
     a__  =merges_file
     a__  =do_lower_case

     # cache of sm.MosesPunctNormalizer instance
     a__  ={}
     # cache of sm.MosesTokenizer instance
     a__  ={}
     a__  ={}

     if langs and len(lowercase_	) == 2:
       a__,				a__  =langs
     else:
       raise ValueError(
           F"""arg `langs` needs to be a list of 2 langs, e.g. ['en', 'ru'], but got {langs}. """
           '''Usually that means that tokenizer can\'t find a mapping for the given model path '''
           '''in PRETRAINED_VOCAB_FILES_MAP, and other maps of this tokenizer.'''	)

     with open(lowercase_,   encoding='''utf-8'''	) as src_vocab_handle:
       a__  =json.load(lowercase_	)
     with open(lowercase_,   encoding='''utf-8'''	) as tgt_vocab_handle:
       a__  =json.load(lowercase_	)
       a__  ={v: k for k, v in tgt_vocab.items()}
     with open(lowercase_,   encoding='''utf-8'''	) as merges_handle:
       a__  =merges_handle.read().split('''\n'''	)[:-1]
     a__  =[tuple(merge.split()[:2]	) for merge in merges]
     a__  =dict(zip(lowercase_,   range(len(lowercase_	)	)	)	)
     a__  ={}



   def 			_UpperCAmelCase			(    self	)							->   Dict[str, int]:






     """simple docstring"""






     return self.get_src_vocab()



   @property
   def 			_UpperCAmelCase			(    self	)							->   int:






     """simple docstring"""






     return self.src_vocab_size



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_	)							->   Optional[Any]:






     """simple docstring"""






     if lang not in self.cache_moses_punct_normalizer:
       a__  =self.sm.MosesPunctNormalizer(lang=lowercase_	)
       a__  =punct_normalizer
     return self.cache_moses_punct_normalizer[lang].normalize(lowercase_	)



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_	)							->   List[str]:






     """simple docstring"""






     if lang not in self.cache_moses_tokenizer:
       a__  =self.sm.MosesTokenizer(lang=lowercase_	)
       a__  =moses_tokenizer
     return self.cache_moses_tokenizer[lang].tokenize(
         lowercase_,   aggressive_dash_splits=lowercase_,   return_str=lowercase_,   escape=lowercase_	)



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_	)							->   Dict:






     """simple docstring"""






     if lang not in self.cache_moses_detokenizer:
       a__  =self.sm.MosesDetokenizer(lang=lowercase_	)
       a__  =moses_detokenizer
     return self.cache_moses_detokenizer[lang].detokenize(lowercase_	)



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






     """simple docstring"""






     a__  =replace_unicode_punct(lowercase_	)
     a__  =self.moses_punct_norm(lowercase_,   lowercase_	)
     a__  =remove_non_printing_char(lowercase_	)
     return text



   @property
   def 			_UpperCAmelCase			(    self	)							->   List[str]:






     """simple docstring"""






     return len(self.encoder	)



   @property
   def 			_UpperCAmelCase			(    self	)							->   List[str]:






     """simple docstring"""






     return len(self.decoder	)



   def 			_UpperCAmelCase			(    self	)							->   Any:






     """simple docstring"""






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



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






     """simple docstring"""






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



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






     """simple docstring"""






     a__  =tuple(token[:-1]	) + (token[-1] + '''</w>''',)
     if token in self.cache:
       return self.cache[token]
     a__  =get_pairs(lowercase_	)

     if not pairs:
       return token + "</w>"

     while True:
       a__  =min(lowercase_,   key=lambda lowercase_	: self.bpe_ranks.get(lowercase_,   float('''inf'''	)	)	)
       if bigram not in self.bpe_ranks:
         break
       a__,				a__  =bigram
       a__  =[]
       a__  =0
       while i < len(lowercase_	):
         try:
           a__  =word.index(lowercase_,   lowercase_	)
         except ValueError:
           new_word.extend(word[i:]	)
           break
         else:
           new_word.extend(word[i:j]	)
           a__  =j

         if word[i] == first and i < len(lowercase_	) - 1 and word[i + 1] == second:
           new_word.append(first + second	)
           i += 2
         else:
           new_word.append(word[i]	)
           i += 1
       a__  =tuple(lowercase_	)
       a__  =new_word
       if len(lowercase_	) == 1:
         break
       else:
         a__  =get_pairs(lowercase_	)
     a__  =''' '''.join(lowercase_	)
     if word == "\n  </w>":
       a__  ='''\n</w>'''
     a__  =word
     return word



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_="en",   lowercase_=False	)							->   Dict:






     """simple docstring"""






     a__  =self.src_lang

     if self.do_lower_case:
       a__  =text.lower()

     if bypass_tokenizer:
       a__  =text.split()
     else:
       a__  =self.moses_pipeline(lowercase_,   lang=lowercase_	)
       a__  =self.moses_tokenize(lowercase_,   lang=lowercase_	)

     a__  =[]
     for token in text:
       if token:
         split_tokens.extend(list(self.bpe(lowercase_	).split(''' '''	)	)	)

     return split_tokens



   def 			_UpperCAmelCase			(    self,   lowercase_	)							->   int:






     """simple docstring"""






     return self.encoder.get(lowercase_,   self.encoder.get(self.unk_token	)	)



   def 			_UpperCAmelCase			(    self,   lowercase_	)							->   List[str]:






     """simple docstring"""






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



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






     """simple docstring"""






     a__  =[t.replace(''' ''',   ''''''	).replace('''</w>''',   ''' '''	) for t in tokens]
     a__  =''''''.join(lowercase_	).split()
     # detokenize
     a__  =self.moses_detokenize(lowercase_,   self.tgt_lang	)
     return text



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_ = None	)							->   List[int]:






     """simple docstring"""






     a__  =[self.sep_token_id]

     # no bos used in fairseq
     if token_ids_a is None:
       return token_ids_a + sep
     return token_ids_a + sep + token_ids_a + sep



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_ = None,   lowercase_ = False	)							->   List[int]:






     """simple docstring"""






     if already_has_special_tokens:
       return super().get_special_tokens_mask(
           token_ids_a=lowercase_,   token_ids_a=lowercase_,   already_has_special_tokens=lowercase_	)
     # no bos used in fairseq
     if token_ids_a is not None:
       return ([0] * len(lowercase_	)) + [1] + ([0] * len(lowercase_	)) + [1]
     return ([0] * len(lowercase_	)) + [1]



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_ = None	)							->   List[int]:






     """simple docstring"""






     a__  =[self.sep_token_id]

     # no bos used in fairseq
     if token_ids_a is None:
       return len(token_ids_a + sep	) * [0]
     return len(token_ids_a + sep	) * [0] + len(token_ids_a + sep	) * [1]



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_ = None	)							->   Tuple[str]:






     """simple docstring"""






     if not os.path.isdir(lowercase_	):
       logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""	)
       return

     a__  =os.path.join(
         lowercase_,   (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''src_vocab_file''']	)
     a__  =os.path.join(
         lowercase_,   (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''tgt_vocab_file''']	)
     a__  =os.path.join(
         lowercase_,   (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']	)

     with open(lowercase_,   '''w''',   encoding='''utf-8'''	) as f:
       f.write(json.dumps(self.encoder,   indent=2,   sort_keys=lowercase_,   ensure_ascii=lowercase_	) + '''\n'''	)

     with open(lowercase_,   '''w''',   encoding='''utf-8'''	) as f:
       a__  ={v: k for k, v in self.decoder.items()}
       f.write(json.dumps(lowercase_,   indent=2,   sort_keys=lowercase_,   ensure_ascii=lowercase_	) + '''\n'''	)

     a__  =0
     with open(lowercase_,   '''w''',   encoding='''utf-8'''	) as writer:
       for bpe_tokens, token_index in sorted(self.bpe_ranks.items(),   key=lambda lowercase_	: kv[1]	):
         if index != token_index:
           logger.warning(
               F"""Saving vocabulary to {merges_file}: BPE merge indices are not consecutive."""
               ''' Please check that the tokenizer is not corrupted!'''	)
           a__  =token_index
         writer.write(''' '''.join(lowercase_	) + '''\n'''	)
         index += 1

     return src_vocab_file, tgt_vocab_file, merges_file



   def __getstate__(    self	)							->   int:






     """simple docstring"""






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



   def __setstate__(    self,   lowercase_	)							->   str:






     """simple docstring"""






     a__  =d

     try:
       import sacremoses
     except ImportError:
       raise ImportError(
           '''You need to install sacremoses to use XLMTokenizer. '''
           '''See https://pypi.org/project/sacremoses/ for installation.'''	)

     a__  =sacremoses



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
       import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


lowerCamelCase									= get_logger(__name__)




def       UpperCAmelCase__    (      _A				:			str     ,				_A				:			List[dict]     ,				_A				:			list     ,				_A				:			list     ,				_A				:			int     ,				_A				:			GradScaler = None     ,				):





  '''simple docstring'''







  for i, state in enumerate(_A							):
    a__  =F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Model weights saved in {output_model_file}"""							)
  # Optimizer states
  for i, opt in enumerate(_A							):
    a__  =opt.state_dict()
    a__  =F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Optimizer state saved in {output_optimizer_file}"""							)
  # Scheduler states
  for i, scheduler in enumerate(_A							):
    a__  =scheduler.state_dict()
    a__  =F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Scheduler state saved in {output_scheduler_file}"""							)
  # GradScaler state
  if scaler is not None:
    a__  =scaler.state_dict()
    a__  =os.path.join(_A     ,				_A							)
    torch.save(_A     ,				_A							)
    logger.info(F"""Gradient scaler state saved in {output_scaler_file}"""							)
  # Random number generator states
  a__  ={}
  a__  =F"""{RNG_STATE_NAME}_{process_index}.pkl"""
  a__  =random.getstate()
  a__  =np.random.get_state()
  a__  =torch.get_rng_state()
  if is_xpu_available():
    a__  =torch.xpu.get_rng_state_all()
  else:
    a__  =torch.cuda.get_rng_state_all()
  if is_tpu_available():
    a__  =xm.get_rng_state()
  a__  =os.path.join(_A     ,				_A							)
  torch.save(_A     ,				_A							)
  logger.info(F"""Random states saved in {output_states_file}"""							)
  return output_dir




def       UpperCAmelCase__    (      _A				:			Tuple     ,				_A				:			List[Any]     ,				_A				:			int     ,				_A				:			Tuple     ,				_A				:			List[Any]     ,				_A				:			List[str]=None     ,				_A				:			Dict=None     ,				**_A				:			Tuple     ,				):





  '''simple docstring'''







  if map_location not in [None, "cpu", "on_device"]:
    raise TypeError(
        '''Unsupported optimizer map location passed, please choose one of `None`, `\'cpu\'`, or `\'on_device\'`'''							)
  if map_location is None:
    a__  ='''cpu'''
  elif map_location == "on_device":
    a__  =PartialState().device
  # Model states
  for i, model in enumerate(_A							):
    a__  =F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    models[i].load_state_dict(torch.load(_A     ,				map_location=_A							)     ,				**_A							)
  logger.info('''All model weights loaded successfully'''							)

  # Optimizer states
  for i, opt in enumerate(_A							):
    a__  =F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    a__  =torch.load(_A     ,				map_location=_A							)
    optimizers[i].load_state_dict(_A							)
  logger.info('''All optimizer states loaded successfully'''							)

  # Scheduler states
  for i, scheduler in enumerate(_A							):
    a__  =F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    scheduler.load_state_dict(torch.load(_A							)							)
  logger.info('''All scheduler states loaded successfully'''							)

  # GradScaler state
  if scaler is not None:
    a__  =os.path.join(_A     ,				_A							)
    scaler.load_state_dict(torch.load(_A							)							)
    logger.info('''GradScaler state loaded successfully'''							)

  # Random states
  try:
    a__  =torch.load(os.path.join(_A     ,				F"""{RNG_STATE_NAME}_{process_index}.pkl"""							)							)
    random.setstate(states['''random_state''']							)
    np.random.set_state(states['''numpy_random_seed''']							)
    torch.set_rng_state(states['''torch_manual_seed''']							)
    if is_xpu_available():
      torch.xpu.set_rng_state_all(states['''torch_xpu_manual_seed''']							)
    else:
      torch.cuda.set_rng_state_all(states['''torch_cuda_manual_seed''']							)
    if is_tpu_available():
      xm.set_rng_state(states['''xm_seed''']							)
    logger.info('''All random states loaded successfully'''							)
  except Exception:
    logger.info('''Could not load random states'''							)




def       UpperCAmelCase__    (      _A				:			Tuple     ,				_A				:			Tuple     ,				_A				:			int = 0							):





  '''simple docstring'''







  a__  =Path(_A							) / F"""custom_checkpoint_{index}.pkl"""
  logger.info(F"""Saving the state of {get_pretty_name(_A							)} to {save_location}"""							)
  torch.save(obj.state_dict()     ,				_A							)




def       UpperCAmelCase__    (      _A				:			int     ,				_A				:			Tuple     ,				_A				:			int = 0							):





  '''simple docstring'''







  a__  =F"""{path}/custom_checkpoint_{index}.pkl"""
  logger.info(F"""Loading the state of {get_pretty_name(_A							)} from {load_location}"""							)
  obj.load_state_dict(torch.load(_A     ,				map_location='''cpu'''							)							)
