

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


_UpperCamelCase			       =			logging.get_logger(__name__)

_UpperCamelCase			       =			{
    '''source_spm''': '''source.spm''',
    '''target_spm''': '''target.spm''',
    '''vocab''': '''vocab.json''',
    '''target_vocab_file''': '''target_vocab.json''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
}

_UpperCamelCase			       =			{
    '''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'''
        )
    },
}

_UpperCamelCase			       =			{'''Helsinki-NLP/opus-mt-en-de''': 512}
_UpperCamelCase			       =			{}



class     lowercase						(       _UpperCamelCase					):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE												=				VOCAB_FILES_NAMES
 __SCREAMING_SNAKE_CASE												=				PRETRAINED_VOCAB_FILES_MAP
 __SCREAMING_SNAKE_CASE												=				PRETRAINED_INIT_CONFIGURATION
 __SCREAMING_SNAKE_CASE												=				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __SCREAMING_SNAKE_CASE												=				["""input_ids""", """attention_mask"""]
 __SCREAMING_SNAKE_CASE												=				re.compile(""">>.+<<"""					)  # type: re.Pattern




 def __init__(self       ,       __a       ,       __a       ,       __a       ,       __a=None       ,       __a=None       ,       __a=None       ,       __a="<unk>"       ,       __a="</s>"       ,       __a="<pad>"       ,       __a=512       ,       __a = None       ,       __a=False       ,       **__a       ,       ) ->					None:






  """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=__a       ,       target_lang=__a       ,       unk_token=__a       ,       eos_token=__a       ,       pad_token=__a       ,       model_max_length=__a       ,       sp_model_kwargs=self.sp_model_kwargs       ,       target_vocab_file=__a       ,       separate_vocabs=__a       ,       **__a       ,       )
  assert Path(__a ).exists(), F"cannot find spm source {source_spm}"

  UpperCAmelCase__							=				separate_vocabs
  UpperCAmelCase__							=				load_json(__a )
  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(__a )
   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(__a       ,       self.sp_model_kwargs )
  UpperCAmelCase__							=				load_spm(__a       ,       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 ) ->					str:






  """simple docstring"""





  try:
   from sacremoses import MosesPunctNormalizer

   UpperCAmelCase__							=				MosesPunctNormalizer(self.source_lang ).normalize
  except (ImportError, FileNotFoundError):
   warnings.warn('Recommended: pip install sacremoses.' )
   UpperCAmelCase__							=				lambda __a : x




 def   UpperCamelCase__							(self       ,       __a ) ->					str:






  """simple docstring"""





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




 def   UpperCamelCase__							(self       ,       __a ) ->					List[Any]:






  """simple docstring"""





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




 def   UpperCamelCase__							(self       ,       __a ) ->					Optional[int]:






  """simple docstring"""





  UpperCAmelCase__							=				self.language_code_re.match(__a )
  UpperCAmelCase__							=				[match.group(0 )] if match else []
  return code, self.language_code_re.sub(''       ,       __a )




 def   UpperCamelCase__							(self       ,       __a ) ->					List[str]:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.remove_language_code(__a )
  UpperCAmelCase__							=				self.current_spm.encode(__a       ,       out_type=__a )
  return code + pieces




 def   UpperCamelCase__							(self       ,       __a ) ->					str:






  """simple docstring"""





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




 def   UpperCamelCase__							(self       ,       __a       ,       **__a ) ->					List[Any]:






  """simple docstring"""





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




 def   UpperCamelCase__							(self       ,       __a       ,       **__a ) ->					Any:






  """simple docstring"""





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




 def   UpperCamelCase__							(self       ,       __a ) ->					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(__a ) + token + " "
    UpperCAmelCase__							=				[]
   else:
    current_sub_tokens.append(__a )
  out_string += sp_model.decode_pieces(__a )
  return out_string.strip()




 def   UpperCamelCase__							(self       ,       __a       ,       __a=None ) ->					List[int]:






  """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 ) ->					Optional[int]:






  """simple docstring"""





  UpperCAmelCase__							=				self.spm_source
  UpperCAmelCase__							=				self.encoder




 def   UpperCamelCase__							(self ) ->					List[Any]:






  """simple docstring"""





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




 @property
 def   UpperCamelCase__							(self ) ->					int:






  """simple docstring"""





  return len(self.encoder )




 def   UpperCamelCase__							(self       ,       __a       ,       __a = None ) ->					Tuple[str]:






  """simple docstring"""





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

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

  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(
       __a       ,       (filename_prefix + '-' if filename_prefix else '') + spm_save_filename )
   if os.path.abspath(__a ) != os.path.abspath(__a ) and os.path.isfile(__a ):
    copyfile(__a       ,       __a )
    saved_files.append(__a )
   elif not os.path.isfile(__a ):
    with open(__a       ,       'wb' ) as fi:
     UpperCAmelCase__							=				spm_model.serialized_model_proto()
     fi.write(__a )
    saved_files.append(__a )

  return tuple(__a )




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  return self.get_src_vocab()




 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





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




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





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




 def __getstate__(self ) ->					Dict:






  """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       ,       __a ) ->					None:






  """simple docstring"""





  UpperCAmelCase__							=				d

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

  UpperCAmelCase__				,    UpperCAmelCase__							=				(load_spm(__a       ,       self.sp_model_kwargs ) for f in self.spm_files)
  UpperCAmelCase__							=				self.spm_source
  self._setup_normalizer()




 def   UpperCamelCase__							(self       ,       *__a       ,       **__a ) ->					Optional[Any]:






  """simple docstring"""





  return 1




 def   UpperCamelCase__							(self       ,       __a ) ->					List[str]:






  """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       ,       __a       ,       __a = None       ,       __a = False ) ->					List[int]:






  """simple docstring"""





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







def 						UpperCamelCase_(			snake_case__:  str  ,    snake_case__:  Dict[str, Any]   )       -> sentencepiece.SentencePieceProcessor:
 UpperCAmelCase__							=				sentencepiece.SentencePieceProcessor(**snake_case__   )
 spm.Load(snake_case__   )
 return spm
def 						UpperCamelCase_(			snake_case__:  str  ,    snake_case__:  str   )       -> None:
 with open(snake_case__  ,    'w'   ) as f:
  json.dump(snake_case__  ,    snake_case__  ,    indent=2   )







def 						UpperCamelCase_(			snake_case__:  str   )       -> Union[Dict, List]:
 with open(snake_case__  ,    'r'   ) as f:
  return json.load(snake_case__   )






import importlib.metadata
import subprocess
import sys





def 	snake_case_      (__A     :						bool = False     )			->	str:
   __lowerCAmelCase  :			Any					    =						False
   if "IPython" in sys.modules:
      __lowerCAmelCase  :			Optional[int]					    =						"""google.colab""" in str(sys.modules["""IPython"""].get_ipython()     )

   if in_colab:
      if upgrade:
         __lowerCAmelCase  :			int					    =						["""pip""", """install""", """-U""", """torch"""]
         subprocess.run(__A					,   check=__A     )
      # get the current version of torch
      __lowerCAmelCase  :			List[Any]					    =						importlib.metadata.version("""torch"""     )
      __lowerCAmelCase  :			List[Any]					    =						torch_version[: torch_version.rindex("""."""     )]
      __lowerCAmelCase  :			int					    =						f'''https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl'''
      __lowerCAmelCase  :			List[Any]					    =						["""pip""", """install""", xla_wheel]
      subprocess.run(__A					,   check=__A     )
   else:
      raise RuntimeError("""`install_xla` utility works only on google colab."""     )


